Ptyrannosaur

Ptyrannosaur#

Ptychography with Robust Artificial Neural Networks Optimized for Sub-Angstrom Accuracy and Ultrafast Reconstruction
arxiv: xxx.xxxx
Kieran Loehr, Rahim Raja,
Xiaochuan Ding, Jeffrey Huang, Gillian M. Nolan, Sang hyun Bae,
Bryan K. Clark and Pinshane Y. Huang

Clark and Huang Research Groups.

This notebook will run Ptyrannosaur,a machine learning model for rapidly reconstructing 2D projected phase information from 4D STEM data, using the models in arxiv:xxx.xxx.xxxx. You can run it on the experimental data used in our paper or on your own data. Not including setup, running the models takes about two minutes on the google colab T4 GPU’s for a 4D STEM dataset of 128x128. (This will be very slow if you are instead running on the CPU so go make sure your Runtime -> Change Runtime Type is set to T4 GPU)

Initial Setup#

You need to run some things to set up the code on colab and download the models. The code in this section needs to be run once (for each colab session you are using) and you can analyze as much data as you want using it.

import sys
import os
import gdown
from IPython.display import Image,display,HTML
from IPython.core.magic import register_cell_magic
!pip install dm_pix -qq

def get_file(drive_path, file_name):
    gdown.download(drive_path,file_name,quiet=False,fuzzy=True)

def get_folder(drive_path):
    gdown.download_folder(url=drive_path, quiet=False, use_cookies=False)

These are the Google drive paths for the PtyRANNOSAUR code, models, and experimental data.

MODEL_DRIVE_PATH = "https://drive.google.com/drive/folders/1p__RvikRSqS0J_Cq5OXP8c3fJ9tVTKOj?usp=share_link"
EXPERIMENTAL_DATA_DRIVE_PATH = "https://drive.google.com/drive/folders/1O6H4c8SUXsfLVwqpTSgehU_FtwQu5VCW?usp=share_link"
PTYRANNOSAUR_DRIVE_PATH = "https://drive.google.com/file/d/1zIlLDVWIA5KZsEcRQl81bqx9GxGSOsS5/view?usp=share_link"

This will pull the code and the models from our google drive.

get_file(PTYRANNOSAUR_DRIVE_PATH, 'ptyran.tar.gz')

get_folder(MODEL_DRIVE_PATH)

Show code cell output

Hide code cell output

Retrieving folder contents

Retrieving folder 1538ZxPZHq7TWx9s27MwMv4qXClTsPwD- Model2
Processing file 1bM7LpCLOAQGoMz4lSrxlO-wA8zLdUnpW experimental_config.json
Processing file 1NNRdkWbDQh6MVBLv1Q-gKC8VzTthuNDK Model2_Final.pkl
Retrieving folder 1wRnzA9Wna6vQWf6kp2QzLyhXHRebpKux Model10
Processing file 1oV9q9cVX4tP0EdFAafeLfCVz6_OCL0M9 experimental_config.json
Processing file 1zyaNO1iO5CaYUGXj3H3UGpf9eVboqxc3 Model10_Final.pkl
Retrieving folder 1d_0KniIA29dbUh9eHAOK3AGG0wOtMYyP Model20
Processing file 1XGixAOuVPSZ57v2MqDJk2jPGfnY4IZvx experimental_config.json
Processing file 1dXuP4nkNjjDdKEs31IW6a6vM3jOT348Y Model20_Final.pkl

Retrieving folder contents completed
Building directory structure
Building directory structure completed
Downloading...
From: https://drive.google.com/uc?id=1bM7LpCLOAQGoMz4lSrxlO-wA8zLdUnpW
To: /content/Models/Model2/experimental_config.json
100%|██████████| 367/367 [00:00<00:00, 659kB/s]
Downloading...
From (original): https://drive.google.com/uc?id=1NNRdkWbDQh6MVBLv1Q-gKC8VzTthuNDK
From (redirected): https://drive.google.com/uc?id=1NNRdkWbDQh6MVBLv1Q-gKC8VzTthuNDK&confirm=t&uuid=e4e37e96-a575-481e-80fd-7b32eefd16c7
To: /content/Models/Model2/Model2_Final.pkl
100%|██████████| 451M/451M [00:06<00:00, 69.7MB/s]
Downloading...
From: https://drive.google.com/uc?id=1oV9q9cVX4tP0EdFAafeLfCVz6_OCL0M9
To: /content/Models/Model10/experimental_config.json
100%|██████████| 369/369 [00:00<00:00, 1.66MB/s]
Downloading...
From (original): https://drive.google.com/uc?id=1zyaNO1iO5CaYUGXj3H3UGpf9eVboqxc3
From (redirected): https://drive.google.com/uc?id=1zyaNO1iO5CaYUGXj3H3UGpf9eVboqxc3&confirm=t&uuid=523a110b-8357-4a59-9fdf-e30004329b90
To: /content/Models/Model10/Model10_Final.pkl
100%|██████████| 451M/451M [00:03<00:00, 143MB/s]
Downloading...
From: https://drive.google.com/uc?id=1XGixAOuVPSZ57v2MqDJk2jPGfnY4IZvx
To: /content/Models/Model20/experimental_config.json
100%|██████████| 369/369 [00:00<00:00, 1.88MB/s]
Downloading...
From (original): https://drive.google.com/uc?id=1dXuP4nkNjjDdKEs31IW6a6vM3jOT348Y
From (redirected): https://drive.google.com/uc?id=1dXuP4nkNjjDdKEs31IW6a6vM3jOT348Y&confirm=t&uuid=dc83066b-fbe7-4201-adba-c00957cd4f05
To: /content/Models/Model20/Model20_Final.pkl
100%|██████████| 451M/451M [00:03<00:00, 141MB/s]
Download completed

This will install the code on colab.

!tar --pax-option='delete=LIBARCHIVE.xattr.*,delete=SCHILY.fflags' --exclude="._*" -xf ptyran.tar.gz
!pip install -e /content/ptyrannosaur
sys.path.insert(0, '/content/ptyrannosaur/src/')

Show code cell output

Hide code cell output

Obtaining file:///content/ptyrannosaur
  Installing build dependencies ... ?25l?25hdone
  Checking if build backend supports build_editable ... ?25l?25hdone
  Getting requirements to build editable ... ?25l?25hdone
  Preparing editable metadata (pyproject.toml) ... ?25l?25hdone
Requirement already satisfied: numpy in /usr/local/lib/python3.12/dist-packages (from ptyrannosaur==0.1.0) (2.0.2)
Requirement already satisfied: scipy in /usr/local/lib/python3.12/dist-packages (from ptyrannosaur==0.1.0) (1.16.3)
Requirement already satisfied: matplotlib in /usr/local/lib/python3.12/dist-packages (from ptyrannosaur==0.1.0) (3.10.0)
Requirement already satisfied: h5py in /usr/local/lib/python3.12/dist-packages (from ptyrannosaur==0.1.0) (3.16.0)
Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib->ptyrannosaur==0.1.0) (1.3.3)
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.12/dist-packages (from matplotlib->ptyrannosaur==0.1.0) (0.12.1)
Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.12/dist-packages (from matplotlib->ptyrannosaur==0.1.0) (4.63.0)
Requirement already satisfied: kiwisolver>=1.3.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib->ptyrannosaur==0.1.0) (1.5.0)
Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.12/dist-packages (from matplotlib->ptyrannosaur==0.1.0) (26.2)
Requirement already satisfied: pillow>=8 in /usr/local/lib/python3.12/dist-packages (from matplotlib->ptyrannosaur==0.1.0) (11.3.0)
Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib->ptyrannosaur==0.1.0) (3.3.2)
Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.12/dist-packages (from matplotlib->ptyrannosaur==0.1.0) (2.9.0.post0)
Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.7->matplotlib->ptyrannosaur==0.1.0) (1.17.0)
Building wheels for collected packages: ptyrannosaur
  Building editable for ptyrannosaur (pyproject.toml) ... ?25l?25hdone
  Created wheel for ptyrannosaur: filename=ptyrannosaur-0.1.0-0.editable-py3-none-any.whl size=2374 sha256=57b4c599e425be99cae3f9f99584c19de7d1e26f4590b61a9d3b5e5558abdc57
  Stored in directory: /tmp/pip-ephem-wheel-cache-zjzx47fl/wheels/02/d6/96/2d7728f400754ba088500310f31bb16f8009db265a863f6755
Successfully built ptyrannosaur
Installing collected packages: ptyrannosaur
Successfully installed ptyrannosaur-0.1.0

Experimental Reconstructions#

Step 1: Getting the data onto colab#

Now you need to get the data into colab. We have currently supplied three models: a 10nm model, a 20nm model, and a 2nm model. See below for the parameters these models are designed for. If you are outside these parameter regions, caveat emptor.

Model	Thickness (nm)	Accelerating Voltage (kV)	Convergence Angle (mrad)	Defocus (nm)	dk (mrad)	Scan Step Size (Å)	Source Size Blur (Å)
Model 2	\(2 \pm 0.6\)	80	25	\(7.5 \pm 2\)	1.67	0.43	\(0.35 \pm 0.15\)
Model 10	\(10 \pm 2\)	300	25	\(5.0 \pm 2\)	0.86	0.5	\(0.35 \pm 0.15\)
Model 20	\(20 \pm 2\)	300	25	\(5.0 \pm 2\)	0.86	0.5	\(0.35 \pm 0.15\)

You can either upload your own data or use our data. In either case the folder structure will look like this (our data will automatically be in this structure).

- Experimental_Data
  - 10nm_Data
    - name_of_your_material (can be anything)
      - scan_x128_y128.raw
    - name_of_another_material (can be anything)
      - scan_x128_y128.raw
  - 20nm_Data
    - name_of_your_material (can be anything)
      - scan_x128_y128.raw
  - 2nm_Data
    - name_of_your_material (can be anything)
      - scan_x128_y128.raw
    - name_of_another_material (can be anything)
      - scan_x128_y128.raw

Please choose one of the three options for getting data

Option 1: Getting our experimental data:

get_folder(EXPERIMENTAL_DATA_DRIVE_PATH)

Option 2: Getting your data from your google drive: (This will be the fastest way of getting your data onto colab)

!mkdir -p /content/Experimental_Data/10nm_Data/
from google.colab import drive
drive.mount('/content/drive')
!cp -r "/content/drive/MyDrive/Your_Drive_Folder_Name" "/content/Experimental_Data/10nm_Data/"

Option 3: Uploading to colab directly from your machine:

Click on the file folder to your left. Then click on the upload button and upload it to the right folder. This will be much slower then uploading to google drive and pulling it from there. You need to be in the correct directory structure (which you can do manually or just pull our data and then use that as a template).

### If you execute this cell, you will get our data (with the correct directory structure)
get_folder(EXPERIMENTAL_DATA_DRIVE_PATH)

Show code cell output

Hide code cell output

Retrieving folder contents

Retrieving folder 1qI4bxZZEEGIkZw689eMjok_4S_9zcW_k 2nm_Data
Retrieving folder 1t3-tNoTyNwO7BA3WKWaQeoyUI76GNAdw 2D Material
Processing file 1Fo6Z0Vnwc3jO8CDDG_RqcLrvSaGaQaT1 scan_x128_y128.raw
Retrieving folder 1SEbq7JCbaAvEBHbVQRGDJLZSLDbBACyh 10nm_Data
Retrieving folder 1pCqzSKNgJQJXOhZZYQMDBZ1iHkTrrwqK Gold Nanoparticle
Processing file 1SDuHfN8Xi_twgdOOn9UzRbACoXW5E0Hu scan_x128_y128.raw
Retrieving folder 1jXcaaxarGDpI1K1WSrmV0DSxq4UqrEb2 Junction
Processing file 1tp0cEZ5YGppa4rY9BWT-bEqIgz7pYj30 scan_x128_y128.raw
Retrieving folder 1ZviLdZVliHC1MTKmM3Q7fTgWLhgFGZny 20nm_Data
Retrieving folder 1AQHHGQ8-LrGJjbAuOw47NIef6K98hChz Sapphire
Processing file 1VnjEpuCpf1s3cvPjf2FxsCm0QnaySbtK scan_x128_y128.raw

Retrieving folder contents completed
Building directory structure
Building directory structure completed
Downloading...
From (original): https://drive.google.com/uc?id=1Fo6Z0Vnwc3jO8CDDG_RqcLrvSaGaQaT1
From (redirected): https://drive.google.com/uc?id=1Fo6Z0Vnwc3jO8CDDG_RqcLrvSaGaQaT1&confirm=t&uuid=28866642-1e53-4968-b19b-4c4b0c32bfed
To: /content/Experimental_Data/2nm_Data/2D Material/scan_x128_y128.raw
100%|██████████| 1.09G/1.09G [00:06<00:00, 171MB/s]
Downloading...
From (original): https://drive.google.com/uc?id=1SDuHfN8Xi_twgdOOn9UzRbACoXW5E0Hu
From (redirected): https://drive.google.com/uc?id=1SDuHfN8Xi_twgdOOn9UzRbACoXW5E0Hu&confirm=t&uuid=2b41c526-dd27-498c-b85e-d602afd61254
To: /content/Experimental_Data/10nm_Data/Gold Nanoparticle/scan_x128_y128.raw
100%|██████████| 1.09G/1.09G [00:13<00:00, 80.4MB/s]
Downloading...
From (original): https://drive.google.com/uc?id=1tp0cEZ5YGppa4rY9BWT-bEqIgz7pYj30
From (redirected): https://drive.google.com/uc?id=1tp0cEZ5YGppa4rY9BWT-bEqIgz7pYj30&confirm=t&uuid=b27cdf16-50cb-4991-aa0a-61d50fe47198
To: /content/Experimental_Data/10nm_Data/Junction/scan_x128_y128.raw
100%|██████████| 1.09G/1.09G [00:12<00:00, 88.6MB/s]
Downloading...
From (original): https://drive.google.com/uc?id=1VnjEpuCpf1s3cvPjf2FxsCm0QnaySbtK
From (redirected): https://drive.google.com/uc?id=1VnjEpuCpf1s3cvPjf2FxsCm0QnaySbtK&confirm=t&uuid=5c107df8-0fda-4849-b967-5dc6d4f261f7
To: /content/Experimental_Data/20nm_Data/Sapphire/scan_x128_y128.raw
100%|██████████| 1.09G/1.09G [00:10<00:00, 104MB/s]
Download completed

Step 2: Doing the reconstructions#

When you run the reconstructions, it will (by default) run reconstructions on all the models in your eval_path using the model in model_path. You can put many folders of materials there at once. When it is done, you can find your images in the folder “Experimental_Evaluation”. To see folders, click on the folder icon to the left (and then sometimes you have to hit the refresh icon)

At this point, it is crucial to be connected to a GPU runtime. CPU runtimes will take a very long time to compute. Even on a T4 GPU the first step will take up to a minute as the model is compiled. Future uses of the model should complete in ~1 minute.

%run /content/ptyrannosaur/scripts/final_model/final_eval_exp.py --model_path '/content/Models/Model2/Model2_Final.pkl' --eval_path '/content/Models/Model2'

Step 3: Viewing the Data#

By default the data is under /content/Experimental_Evaluation/2nm_Data/materialName (you may have to change 2nm_Data to the model you ran the data on. You can access it from the file button on the left or plot it like below. (if you click on the plotted image below it will get larger).

display(Image('/content/Experimental_Data/2nm_Data/2D Material/PtyRANNOSAUR_output.png'))

_images/d9bbfbfeef36e495d9163016265d6929da4abbf725efe0cd0f264ec4eddd7223.png

Step 4: Additional Data#

Here we again reconstruct and plot images for the 10 and 20 nm data (using the respective models).

%run /content/ptyrannosaur/scripts/final_model/final_eval_exp.py --model_path '/content/Models/Model10/Model10_Final.pkl' --eval_path '/content/Models/Model10'

display(Image('/content/Experimental_Data/10nm_Data/Junction/PtyRANNOSAUR_output.png'))

_images/e2e6c02f69731b1d0f3a1ab96df24e66d733db956b4c538fd6f133225152d19d.png

display(Image('/content/Experimental_Data/10nm_Data/Gold Nanoparticle/PtyRANNOSAUR_output.png'))

_images/ba94435550bb57446997925309ab6a038ea06a81174b246accdfad9dfc091fc6.png

%run /content/ptyrannosaur/scripts/final_model/final_eval_exp.py --model_path '/content/Models/Model20/Model20_Final.pkl' --eval_path '/content/Models/Model20'

display(Image('/content/Experimental_Data/20nm_Data/Sapphire/PtyRANNOSAUR_output.png'))

_images/f8ff6e7dfdf5a5815beb5e427937d78f77ec5a3d41dd219f5b0b83ad74ab56f8.png

Pytrannosaur Image Credit: Anna Li

Ptyrannosaur

Contents

Ptyrannosaur#

Initial Setup#

Experimental Reconstructions#

Step 1: Getting the data onto colab#

Step 2: Doing the reconstructions#

Step 3: Viewing the Data#

Step 4: Additional Data#