Running PtyRAD#

Before running this notebook, you must first follow the instruction in README.md to:

  1. Create the Python environment with all dependant Python packages like PyTorch

  2. Activate that python environment

  3. Install ptyrad package into your activated Python environement (only need to install once)

  4. Initialize the ptyrad/ working directory using ptyrad init command

  5. Download and extract the demo data into data/ using the scripts/download_demo_data.py

Note: This notebook is designed for “ease of operation with the solver class PtyRADSolver”, both “reconstruction” and “hypertune” modes are supported.

Author: Chia-Hao Lee, cl2696@cornell.edu

00. Setup Working Directory#

import os

# Change this to the ABSOLUTE PATH to the ptyrad/ folder so you can correctly access data/ and params/
work_dir = "H:/workspace/ptyrad" # Leave this as-is if you're running the notebook from the `ptyrad/scripts/` folder, this will change it back to ptyrad/

os.chdir(work_dir)
print("Current working dir: ", os.getcwd())
# The printed working dir should be ".../ptyrad/" to locate the example params files easily
# Note that the output/ directory will be automatically generated under your working directory

01. Initialize PtyRAD Solver#

from ptyrad.params import load_params
from ptyrad.runtime.device import set_gpu_device
from ptyrad.runtime.diagnostics import print_system_info
from ptyrad.runtime.logging import LoggingManager
from ptyrad.solver import PtyRADSolver

LoggingManager(log_file='ptyrad_log.txt', log_dir='auto', prefix_time='datetime', show_timestamp=True)

# All the following params files are provided in ptyrad/params/ and we're using relative path here
# So if you change the working directory, or have moved params files around, you'll have to provide absolute path to the params file

params_path = "params/examples/tBL_WSe2.yaml"
# params_path = "params/examples/PSO.yaml"
# params_path = "params/walkthrough/00_single_slice.yaml" # You can also explore the entire walkthrough/ folder for step-by-step introduction

print_system_info()

# We enable validation to auto-fill defaults and check parameter consistency since PtyRAD 0.1.0b8
# If you run into issues with validation (e.g., false positives or unexpected errors),
# you can temporarily disable it by setting `validate=False` and prepare a fully complete params file yourself.
# If this happens, please report the bug so we can improve the validation logic.
params = load_params(params_path, validate=True)
device = set_gpu_device(gpuid=0) # Pass in `gpuid = None` if you don't have access to a CUDA-compatible GPU. Note that running PtyRAD with CPU would be much slower than on GPU.

ptycho_solver = PtyRADSolver(params, device=device)

02. Run PtyRAD Solver#

ptycho_solver.run()

# Only `reconstruct` mode will return the final reconstructed model, 
# because it's infeasible to store all models in `hypertune` mode
if not ptycho_solver.if_hypertune:
    model = ptycho_solver.reconstruct_results