cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-10 05:22:41 +02:00
Code Issues Packages Projects Releases Wiki Activity

merge updates to master

Browse Source
This commit is contained in:
levita_a
2025-11-19 10:29:33 +01:00
parent 1dd389d08a 562bcab171
commit ab7048e375
47 changed files with 3206 additions and 1770 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.flake8
+2
View File
@@ -0,0 +1,2 @@
[flake8]
ignore = E501, W503
.github/workflows/main.yml
+46 -12
View File
@@ -7,11 +7,11 @@ on:
branches: [ master ]
jobs:
test:
test-uv-pip-install:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
python-version: ['3.9', '3.10', '3.11', '3.12', '3.13', '3.14']
continue-on-error: true
steps:
@@ -22,14 +22,43 @@ jobs:
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
- name: Install uv
run: |
pip install --upgrade pip
pip install -r requirements.txt
pip install -e . --no-deps
pip install uv
- name: Create virtual environment
run: |
uv venv
- name: Install dependencies with uv
run: |
uv pip install ."[tests]"
- name: Run tests
run: pytest
run: |
uv run pytest
test-pip-install:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ['3.9', '3.14']
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v3
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies with pip
run: |
pip install ."[tests]"
- name: Run tests
run: |
pytest
build-docs:
if: github.event_name == 'push' && github.ref == 'refs/heads/master'
@@ -52,12 +81,17 @@ jobs:
with:
python-version: '3.9'
- name: Install dependencies
- name: Install uv
run: |
pip install --upgrade pip
pip install -r requirements.txt
pip install sphinx sphinx_rtd_theme sphinx-argparse
pip install -e . --no-deps
pip install uv
- name: Create virtual environment
run: |
uv venv
- name: Install dependencies with uv
run: |
uv pip install ."[docs]"
- name: Build docs
working-directory: docs
.github/workflows/publish.yml
+37
View File
@@ -0,0 +1,37 @@
name: Upload Python Package to PyPI when a Release is Created
on:
release:
types: [created]
jobs:
pypi-publish:
name: Publish release to PyPI
runs-on: ubuntu-latest
environment:
name: pypi
url: https://pypi.org/p/cdtools-py
permissions:
id-token: write
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v4
with:
python-version: "3.x"
- name: Install uv
run: |
pip install uv
- name : Create virtual environment
run: |
uv venv
- name: Install project with uv
run: |
uv pip install -e .
- name: Build package with uv
run: |
uv build
- name: Publish package distributions to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
README.md
+50 -1
View File
@@ -26,8 +26,57 @@ model.compare(dataset) # See how the simulated and measured patterns compare
plt.show()
```
Full installation instructions and documentation can be found [here](https://cdtools-developers.github.io/cdtools/).
Further documentation is found [here](https://cdtools-developers.github.io/cdtools/).
# Installation
CDTools can be installed in several ways depending on your needs. For most users, installation from pypi is recommended. For developers or those who want the latest features, installation from source is available.
## Installation from pypi
CDTools can be installed via pip as the [cdtools-py](https://pypi.org/project/cdtools-py/) package on [PyPI](https://pypi.org/):
```bash
$ pip install cdtools-py
```
or using [uv](https://github.com/astral-sh/uv):
```bash
$ uv pip install cdtools-py
```
## Installation from Source
For development or to access the latest features, CDTools can be installed directly from source:
```bash
$ git clone https://github.com/cdtools-developers/cdtools.git
$ cd cdtools
$ pip install -e .
```
or using [uv](https://github.com/astral-sh/uv):
```bash
$ git clone https://github.com/cdtools-developers/cdtools.git
$ cd cdtools
$ uv pip install -e .
```
## Installing for Contributors (with tests and docs dependencies)
If you want to run the test suite or build the documentation, install with the extra dependencies:
```bash
$ pip install -e ."[tests,docs]"
```
or with uv:
```bash
$ uv pip install -e ."[tests,docs]"
```
CDTools was developed in the [photon scattering lab](https://scattering.mit.edu/) at MIT, and further development took place within the [computational x-ray imaging group](https://www.psi.ch/en/cxi) at PSI. The code is distributed under an MIT (a.k.a. Expat) license. If you would like to publish any work that uses CDTools, please contact [Abe Levitan](mailto:abraham.levitan@psi.ch).
docs/source/examples.rst
+8 -2
View File
@@ -31,9 +31,9 @@ When reading this script, note the basic workflow. After the data is loaded, a m
Next, the model is moved to the GPU using the :code:`model.to` function. Any device understood by :code:`torch.Tensor.to` can be specified here. The next line is a bit more subtle - the dataset is told to move patterns to the GPU before passing them to the model using the :code:`dataset.get_as` function. This function does not move the stored patterns to the GPU. If there is sufficient GPU memory, the patterns can also be pre-moved to the GPU using :code:`dataset.to`, but the speedup is empirically quite small.
Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at every epoch, to allow some monitoring code to be run.
Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at the end of every epoch, to allow some monitoring code to be run.
Finally, the results can be studied using :code:`model.inspect(dataet)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset.
Finally, the results can be studied using :code:`model.inspect(dataset)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset.
Fancy Ptycho
@@ -63,6 +63,12 @@ By default, FancyPtycho will also optimize over the following model parameters,
These corrections can be turned off (on) by calling :code:`model.<parameter>.requires_grad = False #(True)`.
Note as well two other changes that are made in this script, when compared to `simple_ptycho.py`. First, a `Reconstructor` object is explicitly created, in this case an `AdamReconstructor`. This object stores a model, dataset, and pytorch optimizer. It is then used to orchestrate the later reconstruction using a call to `Reconstructor.optimize()`.
We use this pattern, instead of the simpler call to `model.Adam_optimize()`, because having the reconstructor store the optimizer as well as the model and dataset allows the moment estimates to persist between multiple rounds of optimization. This leads to the second change: In this script, we run two optimization loops. The first loop aggressively refines the probe, with a low minibatch size and a high learning rate. The second loop has a smaller learning rate and a larger batch size, which allow for a more precise final estimation of the object.
In this case, we used one reconstructor, but it is possible to create additional reconstructors to zero out all the persistant information in the optimizer, if desired, or even to instantiate multiple reconstructors on the same model with different optimization algorithms (e.g. `model.LBFGS_optimize()`).
Gold Ball Ptycho
----------------
docs/source/index.rst
+1
View File
@@ -9,6 +9,7 @@
general
datasets
models
reconstructors
tools/index
indices_tables
docs/source/installation.rst
+84 -47
View File
@@ -1,79 +1,116 @@
Installation
============
Step 1: Download
----------------
CDTools supports python >=3.9 and can be installed via pip as the the `cdtools-py`_ package on `PyPI`_. If you plan to contribute to the code or need a custom environment, installation from source is also possible.
The source code for CDTools is hosted on `Github`_. At the moment, the repository remains private while we decide on licensing. Access can be granted upon request by contacting `Abe Levitan <alevitan@mit.edu>`_.
.. _`cdtools-py`: https://pypi.org/project/cdtools-py/
.. _`PyPI`: https://pypi.org/
.. _`Github`: https://github.com/cdtools-developers/cdtools
Option 1: Installation from PyPI
--------------------------------
The repository remains under active development as of early 2025.
Step 2: Install Dependencies
----------------------------
CDTools is regularly tested with Python versions 3.8 to 3.12, so it is recommended to use one of these versions. In general, CDTools requires Python 3.7 or higher.
The major dependency for CDTools is pytorch (version 1.9.0 or greater). Because the details of the installation can vary depending on platform, GPU availability, etc, it is recommended that you follow the install instructions on `the pytorch site`_ to install pytorch before installing the remaining dependencies.
.. _`the pytorch site`: https://pytorch.org/get-started/locally/
pytorch stopped supporting installation using conda for installation, so it is recommended continue the installation using pip.
To install from `PyPI`_, run:
.. code:: bash
$ pip install -r requirements.txt
$ pip install cdtools-py
This will install all required dependencies and verify that they meet the pytorch version requirements. Additionally, several optional dependencies used for testing and documentation will also be installed. The full set of dependencies and minimum requirements are listed below is listed below.
or you can use `uv`_ for a faster installation:
CDTools is reguarly tested with the latest versions of the packages shown below.
.. _`uv`: https://github.com/astral-sh/uv
Required dependencies:
.. code:: bash
$ uv pip install cdtools-py
Pytorch, a major dependence of CDTools, often needs to be installed with a specific CUDA version for machine compatability. If you run into issues with pytorch, consider first installing pytorch into your environment using the instructions on `the pytorch site`_.
.. _`the pytorch site`: https://pytorch.org/get-started/locally/
Option 2: Installation from source
----------------------------------
Step 1: Download
^^^^^^^^^^^^^^^^
The source code for CDTools is hosted on `Github`_.
.. _`Github`: https://github.com/cdtools-developers/cdtools
To download the source code, you can either clone the repository using git:
.. code:: bash
$ git clone https://github.com/cdtools-developers/cdtools.git
or you can download a zip file of the repository from the `releases page`_.
.. _`releases page`: https://github.com/cdtools-developers/cdtools/releases
Step 2: Install
^^^^^^^^^^^^^^^
Move to the directory where you downloaded the source code. It is recommended that you create a new python virtual environment to install CDTools into.
Installation using pip and uv. Editable mode is recommended for development purposes and is added with the `-e` flag.
.. code:: bash
$ pip install -e .
or using uv:
.. code:: bash
$ uv pip install -e .
To install the required test and documentation dependencies as well, use:
.. code:: bash
$ pip install -e ."[tests,docs]"
or using uv:
.. code:: bash
$ uv pip install -e ."[tests,docs]"
CDTools is reguarly tested with the latest versions of these packages and with python 3.9 through 3.14.
Required dependencies (see pyproject.toml for all details):
* `numpy <http://www.numpy.org>`_ >= 1.0
* `scipy <http://www.scipy.org>`_ >= 1.0
* `matplotlib <https://matplotlib.org>`_ >= 2.0
* `pytorch <https://pytorch.org>`_ >= 1.9.0
* `pytorch <https://pytorch.org>`_ >= 2.3.0
* `python-dateutil <https://github.com/dateutil/dateutil/>`_
* `h5py <https://www.h5py.org/>`_ >= 2.1
Optional dependencies:
Optional dependencies for running tests:
* `pytest <https://docs.pytest.org/>`_
* `pooch <https://www.fatiando.org/pooch/latest/>`_
Optional dependencies for building docs:
* `sphinx <https://www.sphinx-doc.org/>`_ >= 4.3.0
* `sphinx-argparse <https://sphinx-argparse.readthedocs.io>`_
* `sphinx_rtd_theme <https://sphinx-rtd-theme.readthedocs.io/en/stable/>`_ >= 0.5.1
The file "example_environment.yml", included in the repository's top level directory, contains an example of an environment with all dependencies properly installed on a linux machine with a GPU, circa early 2024.
Optional step 4: Run The Tests
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Step 3: Install
---------------
To install CDTools, run the following command from the top level directory (the directory including the setup.py file).
.. code:: bash
$ pip install . --no-deps
This will install a copy of the code, as it exists at the moment of installation. If you would prefer for changes to the code to propagate to the installed version without reinstalling, install the package in developer mode:
.. code:: bash
$ pip install -e . --no-deps
Step 4: Run The Tests
---------------------
To ensure that the installation has worked correctly, it is recommended that you run the unit tests. After ensuring that pytest is installed, run the following command from the top level directory:
To ensure that the installation has worked correctly, it is recommended that you run the unit tests. Execute the following command from the top level directory of the git repository:
.. code:: bash
$ pytest
$ python -m pytest
If any tests fail, make sure that you have all the noted dependencies properly installed. If you do, and things still aren't working, `send me (Abe Levitan) an email <alevitan@mit.edu>`_ and we'll get to the bottom of it. CDTools has been tested on linux and mac, on CPU, CUDA, and MPS.
If any tests fail, make sure that you have all the noted dependencies properly installed. If you do, and things still aren't working, `open an issue on the github page <https://github.com/cdtools-developers/cdtools/issues>`_ and we'll get to the bottom of it.
docs/source/intro.rst
-1
View File
@@ -44,7 +44,6 @@ The high-level interface to CDTools - datasets and models - is built on a set of
- functions for accessing stored data in .cxi files
- plotting tools to visualize data and reconstructions
- basic operations, like light propagators, needed for coherent diffraction
- tools that implement basic operations - such as light propagation - relevant to coherent diffraction.
- analysis functions for assessing the quality of reconstructions
docs/source/reconstructors.rst
+5
View File
@@ -0,0 +1,5 @@
Reconstructors
==============
.. automodule:: cdtools.reconstructors
:members:
example_environment.yml
-181
View File
@@ -1,181 +0,0 @@
name: democdtoolsenv
channels:
- pytorch
- nvidia
- conda-forge
- defaults
dependencies:
- _libgcc_mutex=0.1=main
- _openmp_mutex=5.1=1_gnu
- alabaster=0.7.16=pyhd8ed1ab_0
- babel=2.14.0=pyhd8ed1ab_0
- blas=1.0=mkl
- brotli=1.0.9=h9c3ff4c_4
- bzip2=1.0.8=h7b6447c_0
- c-ares=1.19.1=h5eee18b_0
- ca-certificates=2024.2.2=hbcca054_0
- certifi=2024.2.2=pyhd8ed1ab_0
- charset-normalizer=2.0.4=pyhd3eb1b0_0
- colorama=0.4.6=pyhd8ed1ab_0
- commonmark=0.9.1=py_0
- contourpy=1.2.0=py311hdb19cb5_0
- cuda-cudart=12.1.105=0
- cuda-cupti=12.1.105=0
- cuda-libraries=12.1.0=0
- cuda-nvrtc=12.1.105=0
- cuda-nvtx=12.1.105=0
- cuda-opencl=12.3.101=0
- cuda-runtime=12.1.0=0
- cycler=0.12.1=pyhd8ed1ab_0
- dbus=1.13.18=hb2f20db_0
- docutils=0.20.1=py311h38be061_3
- exceptiongroup=1.2.0=pyhd8ed1ab_2
- expat=2.2.10=h9c3ff4c_0
- ffmpeg=4.3=hf484d3e_0
- filelock=3.13.1=py311h06a4308_0
- fontconfig=2.14.1=hef1e5e3_0
- fonttools=4.25.0=pyhd3eb1b0_0
- freetype=2.12.1=h4a9f257_0
- future=1.0.0=pyhd8ed1ab_0
- glib=2.78.4=h6a678d5_0
- glib-tools=2.78.4=h6a678d5_0
- gmp=6.2.1=h295c915_3
- gmpy2=2.1.2=py311hc9b5ff0_0
- gnutls=3.6.15=he1e5248_0
- gst-plugins-base=1.14.1=h6a678d5_1
- gstreamer=1.14.1=h5eee18b_1
- h5py=3.9.0=py311hdd6beaf_0
- hdf5=1.12.1=h2b7332f_3
- icu=58.2=hf484d3e_1000
- idna=3.4=py311h06a4308_0
- imagesize=1.4.1=pyhd8ed1ab_0
- importlib-metadata=7.0.1=pyha770c72_0
- iniconfig=2.0.0=pyhd8ed1ab_0
- intel-openmp=2023.1.0=hdb19cb5_46306
- jinja2=3.1.3=py311h06a4308_0
- jpeg=9e=h5eee18b_1
- kiwisolver=1.4.4=py311h6a678d5_0
- krb5=1.20.1=h143b758_1
- lame=3.100=h7b6447c_0
- lcms2=2.12=h3be6417_0
- ld_impl_linux-64=2.38=h1181459_1
- lerc=3.0=h295c915_0
- libclang=10.0.1=default_hb85057a_2
- libcublas=12.1.0.26=0
- libcufft=11.0.2.4=0
- libcufile=1.8.1.2=0
- libcurand=10.3.4.107=0
- libcurl=8.5.0=h251f7ec_0
- libcusolver=11.4.4.55=0
- libcusparse=12.0.2.55=0
- libdeflate=1.17=h5eee18b_1
- libedit=3.1.20230828=h5eee18b_0
- libev=4.33=h516909a_1
- libevent=2.1.12=hdbd6064_1
- libffi=3.4.4=h6a678d5_0
- libgcc-ng=11.2.0=h1234567_1
- libgfortran-ng=13.2.0=h69a702a_0
- libgfortran5=13.2.0=ha4646dd_0
- libglib=2.78.4=hdc74915_0
- libgomp=11.2.0=h1234567_1
- libiconv=1.16=h7f8727e_2
- libidn2=2.3.4=h5eee18b_0
- libjpeg-turbo=2.0.0=h9bf148f_0
- libllvm10=10.0.1=he513fc3_3
- libnghttp2=1.57.0=h2d74bed_0
- libnpp=12.0.2.50=0
- libnvjitlink=12.1.105=0
- libnvjpeg=12.1.1.14=0
- libpng=1.6.39=h5eee18b_0
- libpq=12.17=hdbd6064_0
- libssh2=1.10.0=hdbd6064_2
- libstdcxx-ng=11.2.0=h1234567_1
- libtasn1=4.19.0=h5eee18b_0
- libtiff=4.5.1=h6a678d5_0
- libunistring=0.9.10=h27cfd23_0
- libuuid=1.41.5=h5eee18b_0
- libwebp-base=1.3.2=h5eee18b_0
- libxcb=1.15=h7f8727e_0
- libxkbcommon=1.0.1=hfa300c1_0
- libxml2=2.9.14=h74e7548_0
- llvm-openmp=14.0.6=h9e868ea_0
- lz4-c=1.9.4=h6a678d5_0
- markupsafe=2.1.3=py311h5eee18b_0
- matplotlib=3.8.0=py311h06a4308_0
- matplotlib-base=3.8.0=py311ha02d727_0
- mkl=2023.1.0=h213fc3f_46344
- mkl-service=2.4.0=py311h5eee18b_1
- mkl_fft=1.3.8=py311h5eee18b_0
- mkl_random=1.2.4=py311hdb19cb5_0
- mpc=1.1.0=h10f8cd9_1
- mpfr=4.0.2=hb69a4c5_1
- mpmath=1.3.0=py311h06a4308_0
- munkres=1.1.4=pyh9f0ad1d_0
- ncurses=6.4=h6a678d5_0
- nettle=3.7.3=hbbd107a_1
- networkx=3.1=py311h06a4308_0
- nspr=4.35=h6a678d5_0
- nss=3.89.1=h6a678d5_0
- numpy=1.26.3=py311h08b1b3b_0
- numpy-base=1.26.3=py311hf175353_0
- openh264=2.1.1=h4ff587b_0
- openjpeg=2.4.0=h3ad879b_0
- openssl=3.0.13=h7f8727e_0
- packaging=23.2=pyhd8ed1ab_0
- pcre2=10.42=hebb0a14_0
- pillow=10.2.0=py311h5eee18b_0
- pip=23.3.1=py311h06a4308_0
- platformdirs=4.2.0=pyhd8ed1ab_0
- pluggy=1.4.0=pyhd8ed1ab_0
- ply=3.11=py_1
- pooch=1.8.1=pyhd8ed1ab_0
- pygments=2.17.2=pyhd8ed1ab_0
- pyparsing=2.4.7=pyhd8ed1ab_1
- pyqt=5.15.10=py311h6a678d5_0
- pyqt5-sip=12.13.0=py311h5eee18b_0
- pytest=8.0.1=pyhd8ed1ab_1
- python=3.11.7=h955ad1f_0
- python-dateutil=2.8.2=pyhd8ed1ab_0
- python_abi=3.11=2_cp311
- pytorch=2.2.1=py3.11_cuda12.1_cudnn8.9.2_0
- pytorch-cuda=12.1=ha16c6d3_5
- pytorch-mutex=1.0=cuda
- pytz=2024.1=pyhd8ed1ab_0
- pyyaml=6.0.1=py311h5eee18b_0
- qt-main=5.15.2=h327a75a_7
- readline=8.2=h5eee18b_0
- requests=2.31.0=py311h06a4308_1
- scipy=1.11.4=py311h08b1b3b_0
- setuptools=68.2.2=py311h06a4308_0
- sip=6.7.12=py311h6a678d5_0
- six=1.16.0=pyh6c4a22f_0
- snowballstemmer=2.2.0=pyhd8ed1ab_0
- sphinx=7.2.6=pyhd8ed1ab_0
- sphinx-argparse=0.4.0=pyhd8ed1ab_0
- sphinx_rtd_theme=2.0.0=pyha770c72_0
- sphinxcontrib-applehelp=1.0.8=pyhd8ed1ab_0
- sphinxcontrib-devhelp=1.0.6=pyhd8ed1ab_0
- sphinxcontrib-htmlhelp=2.0.5=pyhd8ed1ab_0
- sphinxcontrib-jquery=4.1=pyhd8ed1ab_0
- sphinxcontrib-jsmath=1.0.1=pyhd8ed1ab_0
- sphinxcontrib-qthelp=1.0.7=pyhd8ed1ab_0
- sphinxcontrib-serializinghtml=1.1.10=pyhd8ed1ab_0
- sqlite=3.41.2=h5eee18b_0
- sympy=1.12=py311h06a4308_0
- tbb=2021.8.0=hdb19cb5_0
- tk=8.6.12=h1ccaba5_0
- tomli=2.0.1=pyhd8ed1ab_0
- torchaudio=2.2.1=py311_cu121
- torchtriton=2.2.0=py311
- torchvision=0.17.1=py311_cu121
- tornado=6.3.3=py311h5eee18b_0
- typing_extensions=4.9.0=py311h06a4308_1
- tzdata=2023d=h04d1e81_0
- urllib3=2.1.0=py311h06a4308_0
- wheel=0.41.2=py311h06a4308_0
- xz=5.4.5=h5eee18b_0
- yaml=0.2.5=h7b6447c_0
- zipp=3.17.0=pyhd8ed1ab_0
- zlib=1.2.13=h5eee18b_0
- zstd=1.5.5=hc292b87_0
examples/fancy_ptycho.py
+12 -4
View File
@@ -19,19 +19,27 @@ device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
# For this script, we use a slightly different pattern where we explicitly
# create a `Reconstructor` class to orchestrate the reconstruction. The
# reconstructor will store the model and dataset and create an appropriate
# optimizer. This allows the optimizer to persist between loops, along with
# e.g. estimates of the moments of individual parameters
recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
# The learning rate parameter sets the alpha for Adam.
# The beta parameters are (0.9, 0.999) by default
# The batch size sets the minibatch size
for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
for loss in recon.optimize(50, lr=0.02, batch_size=10):
print(model.report())
# Plotting is expensive, so we only do it every tenth epoch
if model.epoch % 10 == 0:
model.inspect(dataset)
# It's common to chain several different reconstruction loops. Here, we
# started with an aggressive refinement to find the probe, and now we
# polish the reconstruction with a lower learning rate and larger minibatch
for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50):
# started with an aggressive refinement to find the probe in the previous
# loop, and now we polish the reconstruction with a lower learning rate
# and larger minibatch
for loss in recon.optimize(50, lr=0.005, batch_size=50):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
examples/gold_ball_ptycho.py
+7 -4
View File
@@ -29,6 +29,7 @@ model = cdtools.models.FancyPtycho.from_dataset(
probe_fourier_crop=pad
)
# This is a trick that my grandmother taught me, to combat the raster grid
# pathology: we randomze the our initial guess of the probe positions.
# The units here are pixels in the object array.
@@ -42,17 +43,20 @@ device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
# Create the reconstructor
recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
# This will save out the intermediate results if an exception is thrown
# during the reconstruction
with model.save_on_exception(
'example_reconstructions/gold_balls_earlyexit.h5', dataset):
for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
for loss in recon.optimize(20, lr=0.005, batch_size=50):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100):
for loss in recon.optimize(50, lr=0.002, batch_size=100):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
@@ -64,8 +68,7 @@ with model.save_on_exception(
# Setting schedule=True automatically lowers the learning rate if
# the loss fails to improve after 10 epochs
for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100,
schedule=True):
for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
examples/gold_ball_split.py
+6 -3
View File
@@ -36,15 +36,18 @@ for label, dataset in zip(labels, datasets):
model.to(device=device)
dataset.get_as(device=device)
# Create the reconstructor
recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
# For batched reconstructions like this, there's no need to live-plot
# the progress
for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
for loss in recon.optimize(20, lr=0.005, batch_size=50):
print(model.report())
for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100):
for loss in recon.optimize(50, lr=0.002, batch_size=100):
print(model.report())
for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100,
for loss in recon.optimize(100, lr=0.001, batch_size=100,
schedule=True):
print(model.report())
examples/simple_ptycho.py
+10 -1
View File
@@ -1,3 +1,12 @@
"""
Runs a very simple reconstruction using the SimplePtycho model, which was
designed to be an easy introduction to show how the models are made and used.
For a more realistic example of how to use cdtools for real-world data,
look at fancy_ptycho.py and gold_ball_ptycho.py, both of which use the
more powerful FancyPtycho model and include more information on how to
correct for common sources of error.
"""
import cdtools
from matplotlib import pyplot as plt
@@ -13,7 +22,7 @@ device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
# We run the actual reconstruction
# We run the reconstruction
for loss in model.Adam_optimize(100, dataset, batch_size=10):
# We print a quick report of the optimization status
print(model.report())
pyproject.toml
+51 -1
View File
@@ -1,3 +1,53 @@
[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "cdtools-py"
description = "Tools for coherent diffractive imaging and ptychography"
readme = "README.md"
requires-python = ">=3.9"
license = { file = "LICENSE.txt" }
authors = [
{ name = "Abe Levitan", email = "abraham.levitan@psi.ch" },
{ name = "Dayne Y. Sasaki" },
{ name = "Damian Guenzing" },
{ name = "Madelyn Cain" },
{ name = "Anastasiia Kutakh" }
]
maintainers = [
{ name = "Abe Levitan", email = "abraham.levitan@psi.ch" }
]
keywords = ["ptychography", "CDI", "imaging", "torch", "differentiable"]
classifiers = [
"Programming Language :: Python :: 3",
"Operating System :: OS Independent",
"License :: OSI Approved :: MIT License"
]
urls = { "Homepage" = "https://github.com/cdtools-developers/cdtools", "Documentation" = "https://cdtools-developers.github.io/cdtools/" }
dependencies = [
"numpy>=1.0",
"scipy>=1.0",
"matplotlib>=2.0",
"torch>=2.3.0",
"h5py>=2.1",
"python-dateutil",
]
dynamic = ["version"]
[tool.setuptools.dynamic]
version = {attr = "cdtools._version.__version__"}
[project.optional-dependencies]
tests = [
"pytest",
"pooch"
]
docs = [
"sphinx>=4.3.0",
"sphinx-argparse",
"sphinx_rtd_theme>=0.5.1"
]
[tool.ruff]
# Decrease the maximum line length to 79 characters.
line-length = 79
requirements.txt
-11
View File
@@ -1,11 +0,0 @@
numpy>=1.0
scipy>=1.0
matplotlib>=2.0 # 2.0 introduces better colormaps which are used by default
torch>=1.9.0 #1.9.0 implements support for autograd on indexed complex tensors
h5py>=2.1
python-dateutil
pytest
pooch
sphinx>=4.3.0 # Fixes a bug with bulleted lists
sphinx-argparse
sphinx_rtd_theme>=0.5.1 # Fixes a bug with bulleted lists
setup.py
-41
View File
@@ -1,41 +0,0 @@
import setuptools
with open("README.md", "r") as fh:
long_description = fh.read()
setuptools.setup(
name="cdtools",
version="0.2.0",
python_requires='>3.8', # recommended minimum version for pytorch 2.3.0
author="Abe Levitan",
author_email="abraham.levitan@psi.ch",
description="Tools for coherent diffractive imaging and ptychography",
long_description=long_description,
long_description_content_type="text/markdown",
url="https://github.mit.edu/scattering/CDTools.git",
install_requires=[
"numpy>=1.0",
"scipy>=1.0",
"matplotlib>=2.0", # 2.0 has better colormaps which are used by default
"python-dateutil",
"torch>=2.3.0", #2.3.0 is the earliest release for which L-BFGS works directly on complex-valued leaf tensors
"h5py>=2.1"],
extras_require={
'tests': [
"pytest",
"pooch",
],
'docs': [
"sphinx>=4.3.0",
"sphinx-argparse",
"sphinx_rtd_theme>=0.5.1"
]
},
package_dir={"": "src"},
packages=setuptools.find_packages("src"),
classifiers=[
"Programming Language :: Python :: 3",
"Operating System :: OS Independent",
],
)
src/cdtools/__init__.py
+4 -2
View File
@@ -4,9 +4,11 @@ import warnings
warnings.filterwarnings("ignore",
message='To copy construct from a tensor, ')
__all__ = ['tools', 'datasets', 'models']
__all__ = ['tools', 'datasets', 'models', 'reconstructors']
from ._version import __version__
from cdtools import tools
from cdtools import datasets
from cdtools import models
from cdtools import reconstructors
src/cdtools/_version.py
+1
View File
@@ -0,0 +1 @@
__version__ = "0.3.1.dev"
src/cdtools/datasets/base.py
+30 -9
View File
@@ -34,10 +34,16 @@ class CDataset(torchdata.Dataset):
needed to allow for easy mixing of data on the CPU and GPU.
"""
def __init__(self, entry_info=None, sample_info=None,
wavelength=None,
detector_geometry=None, mask=None,
background=None):
def __init__(
self,
entry_info=None,
sample_info=None,
wavelength=None,
detector_geometry=None,
mask=None,
qe_mask=None,
background=None,
):
"""The __init__ function allows construction from python objects.
@@ -73,6 +79,12 @@ class CDataset(torchdata.Dataset):
self.mask = t.tensor(mask, dtype=t.bool)
else:
self.mask = None
if qe_mask is not None:
self.qe_mask = t.as_tensor(qe_mask, dtype=t.float32)
else:
self.qe_mask = None
if background is not None:
self.background = t.tensor(background, dtype=t.float32)
else:
@@ -98,6 +110,8 @@ class CDataset(torchdata.Dataset):
if self.mask is not None:
self.mask = self.mask.to(*args,**mask_kwargs)
if self.qe_mask is not None:
self.qe_mask = self.qe_mask.to(*args,**kwargs)
if self.background is not None:
self.background = self.background.to(*args,**kwargs)
@@ -193,12 +207,17 @@ class CDataset(torchdata.Dataset):
'basis' : basis,
'corner' : corner}
mask = cdtdata.get_mask(cxi_file)
qe_mask = cdtdata.get_qe_mask(cxi_file)
dark = cdtdata.get_dark(cxi_file)
return cls(entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask, background=dark)
return cls(
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask,
qe_mask=qe_mask,
background=dark,
)
def to_cxi(self, cxi_file):
@@ -236,6 +255,8 @@ class CDataset(torchdata.Dataset):
corner = corner)
if self.mask is not None:
cdtdata.add_mask(cxi_file, self.mask)
if self.qe_mask is not None:
cdtdata.add_qe_mask(cxi_file, self.qe_mask)
if self.background is not None:
cdtdata.add_dark(cxi_file, self.background)
src/cdtools/datasets/ptycho_2d_dataset.py
+66 -24
View File
@@ -1,15 +1,16 @@
import warnings
from copy import copy, deepcopy
import pathlib
import h5py
import numpy as np
import torch as t
from copy import copy
import h5py
import pathlib
from cdtools.datasets import CDataset
from cdtools.datasets.random_selection import random_selection
from cdtools.tools import data as cdtdata
from cdtools.tools import plotting
from matplotlib import pyplot as plt
from cdtools.tools import analysis
from copy import deepcopy
__all__ = ['Ptycho2DDataset']
@@ -77,6 +78,7 @@ class Ptycho2DDataset(CDataset):
self.intensities = t.as_tensor(intensities, dtype=t.float32)
else:
self.intensities = None
def __len__(self):
return self.patterns.shape[0]
@@ -123,9 +125,9 @@ class Ptycho2DDataset(CDataset):
self.translations = self.translations.to(*args, **kwargs)
self.patterns = self.patterns.to(*args, **kwargs)
# It sucks that I can't reuse the base factory method here,
# perhaps there is a way but I couldn't figure it out.
@classmethod
def from_cxi(cls, cxi_file, cut_zeros=True, load_patterns=True):
"""Generates a new Ptycho2DDataset from a .cxi file directly
@@ -147,7 +149,7 @@ class Ptycho2DDataset(CDataset):
"""
# If a bare string is passed
if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
with h5py.File(cxi_file,'r') as f:
with h5py.File(cxi_file, 'r') as f:
return cls.from_cxi(f, cut_zeros=cut_zeros, load_patterns=load_patterns)
# Generate a base dataset
@@ -163,10 +165,15 @@ class Ptycho2DDataset(CDataset):
patterns, axes = cdtdata.get_data(cxi_file, cut_zeros=cut_zeros)
dataset.patterns = t.as_tensor(patterns)
if dataset.patterns.dtype == t.float64:
raise NotImplementedError('64-bit floats are not supported and precision will not be retained in reconstructions! Please explicitly convert your data to 32-bit or submit a pull request')
# If the data is 64-bit, we need to convert it to 32-bit
# because 64-bit floats are not supported in reconstructions
dataset.patterns = dataset.patterns.to(dtype=t.float32)
warnings.warn(
"64-bit floats are not supported and precision will not be retained in reconstructions and were converted to t.float32! "
"If you would like to have 64-bit support, please open an issue or submit a pull request."
)
dataset.axes = axes
if dataset.mask is None:
dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool)
@@ -175,9 +182,8 @@ class Ptycho2DDataset(CDataset):
dataset.intensities = t.as_tensor(intensities, dtype=t.float32)
except KeyError:
dataset.intensities = None
return dataset
return dataset
def to_cxi(self, cxi_file):
"""Saves out a Ptycho2DDataset as a .cxi file
@@ -269,9 +275,12 @@ class Ptycho2DDataset(CDataset):
"""Plots the mean diffraction pattern across the dataset
The output is normalized so that the summed intensity on the
detector is equal to the total intensity of light that passed
detector is roughly equal to the total intensity of light that passed
through the sample within each detector conjugate field of view.
If the scan points are colinear (which causes issues for this
estimation), the mean pattern is displayed unscaled.
The plot is plotted as log base 10 of the output plus log_offset.
By default, log_offset is set equal to 1, which is a good level for
shot-noise limited data captured in units of photons. More
@@ -372,10 +381,19 @@ class Ptycho2DDataset(CDataset):
equal to the sum of a <factor> x <factor> region of pixels in the
input pattern. This summation is done by pytorch.functional.avg_pool2d.
Any mask and background data which is stored with the dataset is
downsampled with the data. The background is downsampled using the same
method as the data. The mask is expanded so that any output pixel
containing a masked pixel will be masked.
Any mask, quantum efficiency, and background data which is stored with
the dataset is downsampled with the data. The background is downsampled
using the same method as the data.
If there is no quantum efficiency mask, then the mask is downsampled so
that any output pixel containing a masked pixel will be masked. If there
is a quantum efficiency mask, then the quantum efficiency mask is
downsampled using the same method as the data, and the mask is
downsampled to include any pixels for which there is at least one valid
pixel.
To avoid leakage of data from masked pixels, the data is first
multiplied by the mask before downsampling.
Parameters
----------
@@ -383,17 +401,41 @@ class Ptycho2DDataset(CDataset):
Default 2, the factor to downsample by
"""
self.patterns = t.nn.functional.avg_pool2d(
self.patterns.unsqueeze(0), factor, divisor_override=1)[0]
self.mask = t.logical_not(t.nn.functional.max_pool2d(
(1-self.mask.to(dtype=t.uint8)).unsqueeze(0).unsqueeze(0),
factor
)[0,0].to(dtype=t.bool))
if hasattr(self, 'mask') and self.mask is not None:
self.patterns = t.nn.functional.avg_pool2d(
(self.mask * self.patterns).unsqueeze(0),
factor, divisor_override=1)[0]
else:
self.patterns = t.nn.functional.avg_pool2d(
self.patterns.unsqueeze(0),
factor, divisor_override=1)[0]
# If we have a QE mask, we want to include all pixels for which at
# least one of the input pixels was unmasked, because we can account
# for the masked pixels through quantum efficiency
if hasattr(self, 'qe_mask') and self.qe_mask is not None:
self.qe_mask = t.nn.functional.avg_pool2d(
(self.mask * self.qe_mask).unsqueeze(0).unsqueeze(0),
factor)[0,0]
self.mask = t.nn.functional.max_pool2d(
self.mask.to(dtype=t.uint8).unsqueeze(0).unsqueeze(0),
factor)[0,0].to(dtype=t.bool)
# But if there is no QE mask, we need to only preserve pixels for
# which all input pixels were unmasked
elif hasattr(self, 'mask') and self.mask is not None:
self.mask = t.logical_not(t.nn.functional.max_pool2d(
(1-self.mask.to(dtype=t.uint8)).unsqueeze(0).unsqueeze(0),
factor
)[0,0].to(dtype=t.bool))
self.detector_geometry['basis'] = \
self.detector_geometry['basis'] * factor
if self.background is not None:
if hasattr(self, 'background') and self.background is not None:
self.background = t.nn.functional.avg_pool2d(
self.background.unsqueeze(0).unsqueeze(0),
factor,
src/cdtools/models/base.py
+153 -312
View File
@@ -40,6 +40,10 @@ import time
from scipy import io
from contextlib import contextmanager
from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch
from cdtools.reconstructors import AdamReconstructor, LBFGSReconstructor, SGDReconstructor
from cdtools.datasets import CDataset
from typing import List, Union, Tuple
import os
__all__ = ['CDIModel']
@@ -316,202 +320,23 @@ class CDIModel(t.nn.Module):
self.current_checkpoint_id += 1
def AD_optimize(self, iterations, data_loader, optimizer,\
scheduler=None, regularization_factor=None, thread=True,
calculation_width=10):
"""Runs a round of reconstruction using the provided optimizer
This is the basic automatic differentiation reconstruction tool
which all the other, algorithm-specific tools, use. It is a
generator which yields the average loss each epoch, ending after
the specified number of iterations.
By default, the computation will be run in a separate thread. This
is done to enable live plotting with matplotlib during a reconstruction.
If the computation was done in the main thread, this would freeze
the plots. This behavior can be turned off by setting the keyword
argument 'thread' to False.
Parameters
----------
iterations : int
How many epochs of the algorithm to run
data_loader : torch.utils.data.DataLoader
A data loader loading the CDataset to reconstruct
optimizer : torch.optim.Optimizer
The optimizer to run the reconstruction with
scheduler : torch.optim.lr_scheduler._LRScheduler
Optional, a learning rate scheduler to use
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
calculation_width : int
Default 10, how many translations to pass through at once for each round of gradient accumulation. This does not affect the result, but may affect the calculation speed.
Yields
------
loss : float
The summed loss over the latest epoch, divided by the total diffraction pattern intensity
"""
def run_epoch(stop_event=None):
"""Runs one full epoch of the reconstruction."""
# First, initialize some tracking variables
normalization = 0
loss = 0
N = 0
t0 = time.time()
# The data loader is responsible for setting the minibatch
# size, so each set is a minibatch
for inputs, patterns in data_loader:
normalization += t.sum(patterns).cpu().numpy()
N += 1
def closure():
optimizer.zero_grad()
# We further break up the minibatch into a set of chunks.
# This lets us use larger minibatches than can fit
# on the GPU at once, while still doing batch processing
# for efficiency
input_chunks = [[inp[i:i + calculation_width]
for inp in inputs]
for i in range(0, len(inputs[0]),
calculation_width)]
pattern_chunks = [patterns[i:i + calculation_width]
for i in range(0, len(inputs[0]),
calculation_width)]
total_loss = 0
for inp, pats in zip(input_chunks, pattern_chunks):
# This check allows for graceful exit when threading
if stop_event is not None and stop_event.is_set():
exit()
# Run the simulation
sim_patterns = self.forward(*inp)
# Calculate the loss
if hasattr(self, 'mask'):
loss = self.loss(pats,sim_patterns, mask=self.mask)
else:
loss = self.loss(pats,sim_patterns)
# And accumulate the gradients
loss.backward()
total_loss += loss.detach()
# If we have a regularizer, we can calculate it separately,
# and the gradients will add to the minibatch gradient
if regularization_factor is not None \
and hasattr(self, 'regularizer'):
loss = self.regularizer(regularization_factor)
loss.backward()
return total_loss
# This takes the step for this minibatch
loss += optimizer.step(closure).detach().cpu().numpy()
loss /= normalization
# We step the scheduler after the full epoch
if scheduler is not None:
scheduler.step(loss)
self.loss_history.append(loss)
self.epoch = len(self.loss_history)
self.latest_iteration_time = time.time() - t0
self.training_history += self.report() + '\n'
return loss
# We store the current optimizer as a model parameter so that
# it can be saved and loaded for checkpointing
self.current_optimizer = optimizer
# If we don't want to run in a different thread, this is easy
if not thread:
for it in range(iterations):
if self.skip_computation():
self.epoch = self.epoch + 1
if len(self.loss_history) >= 1:
yield self.loss_history[-1]
else:
yield float('nan')
continue
yield run_epoch()
# But if we do want to thread, it's annoying:
else:
# Here we set up the communication with the computation thread
result_queue = queue.Queue()
stop_event = threading.Event()
def target():
try:
result_queue.put(run_epoch(stop_event))
except Exception as e:
# If something bad happens, put the exception into the
# result queue
result_queue.put(e)
# And this actually starts and monitors the thread
for it in range(iterations):
if self.skip_computation():
self.epoch = self.epoch + 1
if len(self.loss_history) >= 1:
yield self.loss_history[-1]
else:
yield float('nan')
continue
calc = threading.Thread(target=target, name='calculator', daemon=True)
try:
calc.start()
while calc.is_alive():
if hasattr(self, 'figs'):
self.figs[0].canvas.start_event_loop(0.01)
else:
calc.join()
except KeyboardInterrupt as e:
stop_event.set()
print('\nAsking execution thread to stop cleanly - please be patient.')
calc.join()
raise e
res = result_queue.get()
# If something went wrong in the thead, we'll get an exception
if isinstance(res, Exception):
raise res
yield res
# And finally, we unset the current optimizer:
self.current_optimizer = None
def Adam_optimize(
self,
iterations,
dataset,
batch_size=15,
lr=0.005,
betas=(0.9, 0.999),
schedule=False,
amsgrad=False,
subset=None,
regularization_factor=None,
iterations: int,
dataset: CDataset,
batch_size: int = 15,
lr: float = 0.005,
betas: Tuple[float] = (0.9, 0.999),
schedule: bool = False,
amsgrad: bool = False,
subset: Union[int, List[int]] = None,
regularization_factor: Union[float, List[float]] = None,
thread=True,
calculation_width=10
):
"""Runs a round of reconstruction using the Adam optimizer
"""
Runs a round of reconstruction using the Adam optimizer from
cdtools.reconstructors.AdamReconstructor.
This is generally accepted to be the most robust algorithm for use
with ptychography. Like all the other optimization routines,
@@ -521,125 +346,132 @@ class CDIModel(t.nn.Module):
Parameters
----------
iterations : int
How many epochs of the algorithm to run
How many epochs of the algorithm to run.
dataset : CDataset
The dataset to reconstruct against
The dataset to reconstruct against.
batch_size : int
Optional, the size of the minibatches to use
Optional, the size of the minibatches to use.
lr : float
Optional, The learning rate (alpha) to use. Defaultis 0.005. 0.05 is typically the highest possible value with any chance of being stable
betas : tuple
Optional, The learning rate (alpha) to use. Defaultis 0.005.
0.05 is typically the highest possible value with any chance
of being stable.
betas : tuple(float)
Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
schedule : float
Optional, whether to use the ReduceLROnPlateau scheduler
schedule : bool
Optional, whether to use the ReduceLROnPlateau scheduler.
amsgrad : bool
Optional, whether to use the AMSGrad variant of this algorithm.
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
regularization_factor : float or list(float).
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each round of gradient accumulation. Does not affect the result, only the calculation speed
Default 10, how many translations to pass through at once for
each round of gradient accumulation. Does not affect the result,
only the calculation speed.
"""
self.training_history += (
f'Planning {iterations} epochs of Adam, with a learning rate = '
f'{lr}, batch size = {batch_size}, regularization_factor = '
f'{regularization_factor}, and schedule = {schedule}.\n'
reconstructor = AdamReconstructor(
model=self,
dataset=dataset,
subset=subset,
)
if subset is not None:
# if subset is just one pattern, turn into a list for convenience
if type(subset) == type(1):
subset = [subset]
dataset = torchdata.Subset(dataset, subset)
# Make a dataloader
data_loader = torchdata.DataLoader(dataset,
batch_size=batch_size,
shuffle=True)
# Define the optimizer
optimizer = t.optim.Adam(
self.parameters(),
lr = lr,
# Run some reconstructions
return reconstructor.optimize(
iterations=iterations,
batch_size=batch_size,
lr=lr,
betas=betas,
amsgrad=amsgrad)
schedule=schedule,
amsgrad=amsgrad,
regularization_factor=regularization_factor, # noqa
thread=thread,
calculation_width=calculation_width,
)
# Define the scheduler
if schedule:
scheduler = t.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2,threshold=1e-9)
else:
scheduler = None
return self.AD_optimize(iterations, data_loader, optimizer,
scheduler=scheduler,
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width)
def LBFGS_optimize(self, iterations, dataset,
lr=0.1,history_size=2, subset=None,
regularization_factor=None, thread=True,
calculation_width=10, line_search_fn=None):
"""Runs a round of reconstruction using the L-BFGS optimizer
def LBFGS_optimize(self,
iterations: int,
dataset: CDataset,
lr: float = 0.1,
history_size: int = 2,
subset: Union[int, List[int]] = None,
regularization_factor: Union[float, List[float]] =None,
thread: bool = True,
calculation_width: int = 10,
line_search_fn: str = None):
"""
Runs a round of reconstruction using the L-BFGS optimizer from
cdtools.reconstructors.LBFGSReconstructor.
This algorithm is often less stable that Adam, however in certain
situations or geometries it can be shockingly efficient. Like all
the other optimization routines, it is defined as a generator
function which yields the average loss each epoch.
Note: There is no batch size, because it is a usually a bad idea to use
NOTE: There is no batch size, because it is a usually a bad idea to use
LBFGS on anything but all the data at onece
Parameters
----------
iterations : int
How many epochs of the algorithm to run
How many epochs of the algorithm to run.
dataset : CDataset
The dataset to reconstruct against
The dataset to reconstruct against.
lr : float
Optional, the learning rate to use
Optional, the learning rate to use.
history_size : int
Optional, the length of the history to use.
subset : list(int) or int
Optional, a pattern index or list of pattern indices to ues
Optional, a pattern index or list of pattern indices to use.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation.
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. Does not affect the result, only
the calculation speed.
"""
if subset is not None:
# if just one pattern, turn into a list for convenience
if type(subset) == type(1):
subset = [subset]
dataset = torchdata.Subset(dataset, subset)
# Make a dataloader. This basically does nothing but load all the
# data at once
data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset))
# Define the optimizer
optimizer = t.optim.LBFGS(self.parameters(),
lr = lr, history_size=history_size,
line_search_fn=line_search_fn)
return self.AD_optimize(iterations, data_loader, optimizer,
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width)
def SGD_optimize(self, iterations, dataset, batch_size=None,
lr=0.01, momentum=0, dampening=0, weight_decay=0,
nesterov=False, subset=None, regularization_factor=None,
thread=True, calculation_width=10):
"""Runs a round of reconstruction using the SGD optimizer
reconstructor = LBFGSReconstructor(
model=self,
dataset=dataset,
subset=subset,
)
# Run some reconstructions
return reconstructor.optimize(
iterations=iterations,
lr=lr,
history_size=history_size,
regularization_factor=regularization_factor, # noqa
thread=thread,
calculation_width=calculation_width,
line_search_fn=line_search_fn,
)
def SGD_optimize(self,
iterations: int,
dataset: CDataset,
batch_size: int = None,
lr: float = 2e-7,
momentum: float = 0,
dampening: float = 0,
weight_decay: float = 0,
nesterov: bool = False,
subset: Union[int, List[int]] = None,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10):
"""
Runs a round of reconstruction using the SGD optimizer from
cdtools.reconstructors.SGDReconstructor.
This algorithm is often less stable that Adam, but it is simpler
and is the basic workhorse of gradience descent.
@@ -647,51 +479,54 @@ class CDIModel(t.nn.Module):
Parameters
----------
iterations : int
How many epochs of the algorithm to run
How many epochs of the algorithm to run.
dataset : CDataset
The dataset to reconstruct against
The dataset to reconstruct against.
batch_size : int
Optional, the size of the minibatches to use
Optional, the size of the minibatches to use.
lr : float
Optional, the learning rate to use
Optional, the learning rate to use.
momentum : float
Optional, the length of the history to use.
dampening : float
Optional, dampening for the momentum.
weight_decay : float
Optional, weight decay (L2 penalty).
nesterov : bool
Optional, enables Nesterov momentum. Only applicable when momentum
is non-zero.
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use
Optional, a pattern index or list of pattern indices to use.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 1, how many translations to pass through at once for each round of gradient accumulation
Default 10, how many translations to pass through at once for each
round of gradient accumulation.
"""
reconstructor = SGDReconstructor(
model=self,
dataset=dataset,
subset=subset,
)
if subset is not None:
# if just one pattern, turn into a list for convenience
if type(subset) == type(1):
subset = [subset]
dataset = torchdata.Subset(dataset, subset)
# Make a dataloader
if batch_size is not None:
data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
shuffle=True)
else:
data_loader = torchdata.DataLoader(dataset)
# Define the optimizer
optimizer = t.optim.SGD(self.parameters(),
lr = lr, momentum=momentum,
dampening=dampening,
weight_decay=weight_decay,
nesterov=nesterov)
return self.AD_optimize(iterations, data_loader, optimizer,
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width)
# Run some reconstructions
return reconstructor.optimize(
iterations=iterations,
batch_size=batch_size,
lr=lr,
momentum=momentum,
dampening=dampening,
weight_decay=weight_decay,
nesterov=nesterov,
regularization_factor=regularization_factor, # noqa
thread=thread,
calculation_width=calculation_width,
)
def report(self):
@@ -872,7 +707,13 @@ class CDIModel(t.nn.Module):
updating = True if len(axes[0].images) >= 1 else False
inputs, output = dataset[idx:idx+1]
sim_data = self.forward(*inputs).detach().cpu().numpy()[0]
sim_data = self.forward(*inputs).detach().cpu().numpy()
# The length of sim_data.shape changes when you're doing
# either a ptycho (3) or an RPI (2) reconstruction.
# We need to make sure that sim_data is 2D.
if len(sim_data.shape) > 2:
sim_data = sim_data[0]
meas_data = output.detach().cpu().numpy()[0]
if hasattr(self, 'mask') and self.mask is not None:
mask = self.mask.detach().cpu().numpy()
src/cdtools/models/fancy_ptycho.py
+83 -9
View File
@@ -28,6 +28,7 @@ class FancyPtycho(CDIModel):
probe_fourier_shifts=None,
mask=None,
weights=None,
qe_mask=None,
translation_scale=1,
saturation=None,
probe_support=None,
@@ -109,7 +110,18 @@ class FancyPtycho(CDIModel):
else:
self.register_buffer('mask',
t.as_tensor(mask, dtype=t.bool))
if qe_mask is None:
self.qe_mask = None
else:
self.qe_mask = t.nn.Parameter(
t.as_tensor(qe_mask, dtype=dtype))
# I want the ability to optimize over this, but experience shows
# that it is wildly unstable, so I think it's best to keep
# gradients turned off by default
self.qe_mask.requires_grad=False
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
@@ -224,6 +236,7 @@ class FancyPtycho(CDIModel):
dm_rank=None,
translation_scale=1,
saturation=None,
use_qe_mask=False,
probe_support_radius=None,
probe_fourier_crop=None,
propagation_distance=None,
@@ -454,6 +467,14 @@ class FancyPtycho(CDIModel):
else:
mask = None
if use_qe_mask:
if hasattr(dataset, 'qe_mask') and dataset.qe_mask is not None:
qe_mask = t.as_tensor(dataset.qe_mask, dtype=t.float32)
else:
qe_mask = t.ones(dataset.patterns.shape[-2:], dtype=t.float32)
else:
qe_mask = None
if probe_support_radius is not None:
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
@@ -479,13 +500,15 @@ class FancyPtycho(CDIModel):
weights=Ws,
mask=mask,
background=background,
qe_mask=qe_mask,
translation_scale=translation_scale,
saturation=saturation,
probe_basis=probe_basis,
probe_support=probe_support,
fourier_probe=fourier_probe,
oversampling=oversampling,
loss=loss, units=units,
loss=loss,
units=units,
probe_fourier_shifts=probe_fourier_shifts,
simulate_probe_translation=simulate_probe_translation,
simulate_finite_pixels=simulate_finite_pixels,
@@ -620,6 +643,7 @@ class FancyPtycho(CDIModel):
wavefields,
self.background,
measurement=tools.measurements.incoherent_sum,
qe_mask=self.qe_mask,
saturation=self.saturation,
oversampling=self.oversampling,
simulate_finite_pixels=self.simulate_finite_pixels,
@@ -696,15 +720,40 @@ class FancyPtycho(CDIModel):
def center_probes(self, iterations=4):
"""Centers the probes
"""Centers the probes in real space
Takes the current guess of the illumination function and centers it
using a shift with periodic boundary conditions. It uses
cdtools.tools.image_processing.center internally to do the centering.
Multiple iterations of an algorithm are run, which is helpful if the
illumination is reconstructed near the corners and "wraps around" the
probe field of view.
Note that the centering is always performed in real space, even if
the probe array is defined in Fourier space.
Note that this does not compensate for the centering by adjusting
Note also that this does not compensate for the centering by adjusting
the object, so it's a good idea to reset the object after centering
the probes
Parameters
----------
iterations : int
Default 4, how many iterations of the centering algorithm to run
"""
centered_probe = tools.image_processing.center(
self.probe.data.cpu(), iterations=iterations)
self.probe.data = centered_probe.to(device=self.probe.data.device)
if self.fourier_probe:
prs = tools.propagators.inverse_far_field(self.probe.detach()).cpu()
else:
prs = self.probe.detach().cpu()
centered_prs = tools.image_processing.center(prs, iterations=iterations)
if self.fourier_probe:
self.probe.data = tools.propagators.far_field(
centered_prs.to(device=self.probe.data.device))
else:
self.probe.data = centered_prs.to(device=self.probe.data.device)
def tidy_probes(self):
@@ -840,6 +889,28 @@ class FancyPtycho(CDIModel):
**kwargs),
def plot_translations_and_originals(self, fig, dataset):
"""Only used to make a plot for the plot list."""
p.plot_translations(
dataset.translations,
fig=fig,
units=self.units,
label='original translations',
color='#CCCCCC',
marker='o',
)
p.plot_translations(
self.corrected_translations(dataset),
fig=fig,
units=self.units,
clear_fig=False,
label='refined translations',
color='k',
marker='.'
)
plt.legend()
plot_list = [
('',
lambda self, fig, dataset: self.plot_wavefront_variation(
@@ -937,9 +1008,12 @@ class FancyPtycho(CDIModel):
lambda self: self.exponentiate_obj),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
lambda self, fig, dataset: self.plot_translations_and_originals(fig, dataset)),
('Background',
lambda self, fig: p.plot_amplitude(self.background**2, fig=fig))
lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)),
('Quantum Efficiency Mask',
lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig),
lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None))
]
src/cdtools/reconstructors/__init__.py
+22
View File
@@ -0,0 +1,22 @@
"""
Module `cdtools.tools.reconstructors` contains the `Reconstructor` class and
subclasses which run the ptychography reconstructions on a given model and
dataset.
The reconstructors are designed to resemble so-called 'Trainer' classes that
(in the language of the AI/ML folks) handles the 'training' of a model given
some dataset and optimizer.
"""
# We define __all__ to be sure that import * only imports what we want
__all__ = [
'Reconstructor',
'AdamReconstructor',
'LBFGSReconstructor',
'SGDReconstructor'
]
from cdtools.reconstructors.base import Reconstructor
from cdtools.reconstructors.adam import AdamReconstructor
from cdtools.reconstructors.lbfgs import LBFGSReconstructor
from cdtools.reconstructors.sgd import SGDReconstructor
src/cdtools/reconstructors/adam.py
+178
View File
@@ -0,0 +1,178 @@
"""This module contains the AdamReconstructor subclass for performing
optimization ('reconstructions') on ptychographic/CDI models using
the Adam optimizer.
The Reconstructor class is designed to resemble so-called
'Trainer' classes that (in the language of the AI/ML folks) handles
the 'training' of a model given some dataset and optimizer.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch as t
from typing import Tuple, List, Union
from cdtools.reconstructors import Reconstructor
if TYPE_CHECKING:
from cdtools.models import CDIModel
from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
__all__ = ['AdamReconstructor']
class AdamReconstructor(Reconstructor):
"""
The Adam Reconstructor subclass handles the optimization ('reconstruction')
of ptychographic models and datasets using the Adam optimizer.
Parameters
----------
model: CDIModel
Model for CDI/ptychography reconstruction.
dataset: Ptycho2DDataset
The dataset to reconstruct against.
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use.
schedule : bool
Optional, create a learning rate scheduler
(torch.optim.lr_scheduler._LRScheduler).
Important attributes:
- **model** -- Always points to the core model used.
- **optimizer** -- This class by default uses `torch.optim.Adam` to perform
optimizations.
- **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the
`optimize` method.
- **data_loader** -- A torch.utils.data.DataLoader that is defined by
calling the `setup_dataloader` method.
"""
def __init__(self,
model: CDIModel,
dataset: Ptycho2DDataset,
subset: List[int] = None):
# Define the optimizer for use in this subclass
optimizer = t.optim.Adam(model.parameters())
super().__init__(model, dataset, optimizer, subset=subset)
def adjust_optimizer(self,
lr: int = 0.005,
betas: Tuple[float] = (0.9, 0.999),
amsgrad: bool = False):
"""
Change hyperparameters for the utilized optimizer.
Parameters
----------
lr : float
Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
is typically the highest possible value with any chance of being
stable.
betas : tuple
Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
amsgrad : bool
Optional, whether to use the AMSGrad variant of this algorithm.
"""
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
param_group['betas'] = betas
param_group['amsgrad'] = amsgrad
def optimize(self,
iterations: int,
batch_size: int = 15,
lr: float = 0.005,
betas: Tuple[float] = (0.9, 0.999),
custom_data_loader: t.utils.data.DataLoader = None,
schedule: bool = False,
amsgrad: bool = False,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10,
shuffle: bool = True):
"""
Runs a round of reconstruction using the Adam optimizer
Formerly `CDIModel.Adam_optimize`
This calls the Reconstructor.optimize superclass method
(formerly `CDIModel.AD_optimize`) to run a round of reconstruction
once the dataloader and optimizer hyperparameters have been
set up.
The `batch_size` parameter sets the batch size for the default
dataloader. If a custom data loader is desired, it can be passed
in to the `custom_data_loader` argument, which will override the
`batch_size` and `shuffle` parameters
Parameters
----------
iterations : int
How many epochs of the algorithm to run.
batch_size : int
Optional, the size of the minibatches to use.
lr : float
Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
is typically the highest possible value with any chance of being
stable.
betas : tuple
Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
schedule : bool
Optional, create a learning rate scheduler
(torch.optim.lr_scheduler._LRScheduler).
custom_data_loader : t.utils.data.DataLoader
Optional, a custom DataLoader to use. If set, will override
batch_size.
amsgrad : bool
Optional, whether to use the AMSGrad variant of this algorithm.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. Does not affect the result, only
the calculation speed.
shuffle : bool
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
"""
# Update the training history
self.model.training_history += (
f'Planning {iterations} epochs of Adam, with a learning rate = '
f'{lr}, batch size = {batch_size}, regularization_factor = '
f'{regularization_factor}, and schedule = {schedule}.\n'
)
# The optimizer is created in self.__init__, but the
# hyperparameters need to be set up with self.adjust_optimizer
self.adjust_optimizer(lr=lr,
betas=betas,
amsgrad=amsgrad)
# Set up the scheduler
if schedule:
self.scheduler = \
t.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,
factor=0.2,
threshold=1e-9)
else:
self.scheduler = None
# Now, we run the optimize routine defined in the base class
return super(AdamReconstructor, self).optimize(
iterations,
batch_size=batch_size,
custom_data_loader=custom_data_loader,
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width,
shuffle=shuffle,
)
src/cdtools/reconstructors/base.py
+381
View File
@@ -0,0 +1,381 @@
"""This module contains the base Reconstructor class for performing
optimization ('reconstructions') on ptychographic/CDI models.
The Reconstructor class is designed to resemble so-called
'Trainer' classes that (in the language of the AI/ML folks) handles
the 'training' of a model given some dataset and optimizer.
The subclasses of Reconstructor are required to implement
their own data loaders and optimizer adjusters
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch as t
from torch.utils import data as td
import threading
import queue
import time
from typing import List, Union
if TYPE_CHECKING:
from cdtools.models import CDIModel
from cdtools.datasets import CDataset
__all__ = ['Reconstructor']
class Reconstructor:
"""
Reconstructor handles the optimization ('reconstruction') of ptychographic
models given a CDIModel (or subclass) and corresponding CDataset.
This is a base model that defines all functions Reconstructor subclasses
must implement.
Parameters
----------
model: CDIModel
Model for CDI/ptychography reconstruction
dataset: CDataset
The dataset to reconstruct against
optimizer: torch.optim.Optimizer
The optimizer to use for the reconstruction
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use
Attributes
----------
model : CDIModel
Points to the core model used.
optimizer : torch.optim.Optimizer
Must be defined when initializing the Reconstructor subclass.
scheduler : torch.optim.lr_scheduler, optional
May be defined during the ``optimize`` method.
data_loader : torch.utils.data.DataLoader
Defined by calling the ``setup_dataloader`` method.
"""
def __init__(self,
model: CDIModel,
dataset: CDataset,
optimizer: t.optim.Optimizer,
subset: Union[int, List[int]] = None):
# Store parameters as attributes of Reconstructor
self.model = model
self.optimizer = optimizer
# Store the dataset, clipping it to a subset if needed
if subset is not None:
# if subset is just one pattern, turn into a list for convenience
if isinstance(subset, int):
subset = [subset]
dataset = td.Subset(dataset, subset)
self.dataset = dataset
# Initialize attributes that must be defined by the subclasses
self.scheduler = None
self.data_loader = None
def setup_dataloader(self,
batch_size: int = None,
shuffle: bool = True):
"""
Sets up or re-initializes the dataloader.
Parameters
----------
batch_size : int
Optional, the size of the minibatches to use
shuffle : bool
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
"""
if batch_size is not None:
self.data_loader = td.DataLoader(self.dataset,
batch_size=batch_size,
shuffle=shuffle)
else:
self.data_loader = td.Dataloader(self.dataset)
def adjust_optimizer(self, **kwargs):
"""
Change hyperparameters for the utilized optimizer.
For each optimizer, the keyword arguments should be manually defined
as parameters.
"""
raise NotImplementedError()
def run_epoch(self,
stop_event: threading.Event = None,
regularization_factor: Union[float, List[float]] = None,
calculation_width: int = 10):
"""
Runs one full epoch of the reconstruction. Intended to be called
by Reconstructor.optimize.
Parameters
----------
stop_event : threading.Event
Default None, causes the reconstruction to stop when an exception
occurs in Optimizer.optimize.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. This does not affect the result,
but may affect the calculation speed.
Returns
------
loss : float
The summed loss over the latest epoch, divided by the total
diffraction pattern intensity
"""
# Setting this as an explicit catch makes me feel more comfortable
# exposing it as a public method. This way a user won't be confused
# if they try to use this directly
if self.data_loader is None:
raise RuntimeError(
'No data loader was defined. Please run '
'Reconstructor.setup_dataloader() before running '
'Reconstructor.run_epoch(), or use Reconstructor.optimize(), '
'which does it automatically.'
)
# Initialize some tracking variables
normalization = 0
loss = 0
N = 0
t0 = time.time()
# The data loader is responsible for setting the minibatch
# size, so each set is a minibatch
for inputs, patterns in self.data_loader:
normalization += t.sum(patterns).cpu().numpy()
N += 1
def closure():
self.optimizer.zero_grad()
# We further break up the minibatch into a set of chunks.
# This lets us use larger minibatches than can fit
# on the GPU at once, while still doing batch processing
# for efficiency
input_chunks = [[inp[i:i + calculation_width]
for inp in inputs]
for i in range(0, len(inputs[0]),
calculation_width)]
pattern_chunks = [patterns[i:i + calculation_width]
for i in range(0, len(inputs[0]),
calculation_width)]
total_loss = 0
for inp, pats in zip(input_chunks, pattern_chunks):
# This check allows for graceful exit when threading
if stop_event is not None and stop_event.is_set():
exit()
# Run the simulation
sim_patterns = self.model.forward(*inp)
# Calculate the loss
if hasattr(self.model, 'mask'):
loss = self.model.loss(pats,
sim_patterns,
mask=self.model.mask)
else:
loss = self.model.loss(pats,
sim_patterns)
# And accumulate the gradients
loss.backward()
# Normalize the accumulating total loss
total_loss += loss.detach()
# If we have a regularizer, we can calculate it separately,
# and the gradients will add to the minibatch gradient
if regularization_factor is not None \
and hasattr(self.model, 'regularizer'):
loss = self.model.regularizer(regularization_factor)
loss.backward()
return total_loss
# This takes the step for this minibatch
loss += self.optimizer.step(closure).detach().cpu().numpy()
loss /= normalization
# We step the scheduler after the full epoch
if self.scheduler is not None:
self.scheduler.step(loss)
self.model.loss_history.append(loss)
self.model.epoch = len(self.model.loss_history)
self.model.latest_iteration_time = time.time() - t0
self.model.training_history += self.model.report() + '\n'
return loss
def optimize(self,
iterations: int,
batch_size: int = 1,
custom_data_loader: torch.utils.data.DataLoader = None,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10,
shuffle=True):
"""
Runs a round of reconstruction using the provided optimizer
Formerly CDIModel.AD_optimize
This is the basic automatic differentiation reconstruction tool
which all the other, algorithm-specific tools, use. It is a
generator which yields the average loss each epoch, ending after
the specified number of iterations.
By default, the computation will be run in a separate thread. This
is done to enable live plotting with matplotlib during a
reconstruction.
If the computation was done in the main thread, this would freeze
the plots. This behavior can be turned off by setting the keyword
argument 'thread' to False.
The `batch_size` parameter sets the batch size for the default
dataloader. If a custom data loader is desired, it can be passed
in to the `custom_data_loader` argument, which will override the
`batch_size` and `shuffle` parameters
Please see `AdamReconstructor.optimize()` for an example of how to
override this function when designing a subclass
Parameters
----------
iterations : int
How many epochs of the algorithm to run.
batch_size : int
Optional, the batch size to use. Default is 1. This is typically
overridden by subclasses with an appropriate default for the
specific optimizer.
custom_data_loader : torch.utils.data.DataLoader
Optional, a custom DataLoader to use. Will override batch_size
if set.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. This does not affect the result,
but may affect the calculation speed.
shuffle : bool
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
Yields
------
loss : float
The summed loss over the latest epoch, divided by the total
diffraction pattern intensity.
"""
if custom_data_loader is None:
self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
else:
self.data_loader = custom_data_loader
# We store the current optimizer as a model parameter so that
# it can be saved and loaded for checkpointing
self.current_optimizer = self.optimizer
# If we don't want to run in a different thread, this is easy
if not thread:
for it in range(iterations):
if self.model.skip_computation():
self.epoch = self.epoch + 1
if len(self.model.loss_history) >= 1:
yield self.model.loss_history[-1]
else:
yield float('nan')
continue
yield self.run_epoch(
regularization_factor=regularization_factor, # noqa
calculation_width=calculation_width,
)
# But if we do want to thread, it's annoying:
else:
# Here we set up the communication with the computation thread
result_queue = queue.Queue()
stop_event = threading.Event()
def target():
try:
result_queue.put(
self.run_epoch(
stop_event=stop_event,
regularization_factor=regularization_factor, # noqa
calculation_width=calculation_width,
)
)
except Exception as e:
# If something bad happens, put the exception into the
# result queue
result_queue.put(e)
# And this actually starts and monitors the thread
for it in range(iterations):
if self.model.skip_computation():
self.model.epoch = self.model.epoch + 1
if len(self.model.loss_history) >= 1:
yield self.model.loss_history[-1]
else:
yield float('nan')
continue
calc = threading.Thread(target=target,
name='calculator',
daemon=True)
try:
calc.start()
while calc.is_alive():
if hasattr(self.model, 'figs'):
self.model.figs[0].canvas.start_event_loop(0.01)
else:
calc.join()
except KeyboardInterrupt as e:
stop_event.set()
print('\nAsking execution thread to stop cleanly - ' +
'please be patient.')
calc.join()
raise e
res = result_queue.get()
# If something went wrong in the thead, we'll get an exception
if isinstance(res, Exception):
raise res
yield res
# And finally, we unset the current optimizer:
self.current_optimizer = None
src/cdtools/reconstructors/lbfgs.py
+143
View File
@@ -0,0 +1,143 @@
"""This module contains the LBFGSReconstructor subclass for performing
optimization ('reconstructions') on ptychographic/CDI models using
the LBFGS optimizer.
The Reconstructor class is designed to resemble so-called
'Trainer' classes that (in the language of the AI/ML folks) handles
the 'training' of a model given some dataset and optimizer.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch as t
from typing import List, Union
from cdtools.reconstructors import Reconstructor
if TYPE_CHECKING:
from cdtools.models import CDIModel
from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
__all__ = ['LBFGSReconstructor']
class LBFGSReconstructor(Reconstructor):
"""
The LBFGSReconstructor subclass handles the optimization
('reconstruction') of ptychographic models and datasets using the LBFGS
optimizer.
Parameters
----------
model: CDIModel
Model for CDI/ptychography reconstruction.
dataset: Ptycho2DDataset
The dataset to reconstruct against.
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use.
schedule : bool
Optional, create a learning rate scheduler
(torch.optim.lr_scheduler._LRScheduler).
Important attributes:
- **model** -- Always points to the core model used.
- **optimizer** -- This class by default uses `torch.optim.LBFGS` to
perform optimizations.
- **scheduler** -- A `torch.optim.lr_scheduler` that is defined during
the `optimize` method.
- **data_loader** -- A torch.utils.data.DataLoader that is defined by
calling the `setup_dataloader` method.
"""
def __init__(self,
model: CDIModel,
dataset: Ptycho2DDataset,
subset: List[int] = None):
# Define the optimizer for use in this subclass
optimizer = t.optim.LBFGS(model.parameters())
super().__init__(
model,
dataset,
optimizer,
subset=subset,
)
def adjust_optimizer(self,
lr: int = 0.005,
history_size: int = 2,
line_search_fn: str = None):
"""
Change hyperparameters for the utilized optimizer.
Parameters
----------
lr : float
Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
is typically the highest possible value with any chance of being
stable.
history_size : int
Optional, the length of the history to use.
line_search_fn : str
Optional, either `strong_wolfe` or None
"""
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
param_group['history_size'] = history_size
param_group['line_search_fn'] = line_search_fn
def optimize(self,
iterations: int,
lr: float = 0.1,
history_size: int = 2,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10,
line_search_fn: str = None):
"""
Runs a round of reconstruction using the LBFGS optimizer
Formerly `CDIModel.LBFGS_optimize`
This algorithm is often less stable that Adam, however in certain
situations or geometries it can be shockingly efficient. Like all
the other optimization routines, it is defined as a generator
function which yields the average loss each epoch.
NOTE: There is no batch size, because it is a usually a bad idea to use
LBFGS on anything but all the data at onece
Parameters
----------
iterations : int
How many epochs of the algorithm to run.
lr : float
Optional, The learning rate (alpha) to use. Default is 0.1.
history_size : int
Optional, the length of the history to use.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. Does not affect the result, only
the calculation speed.
"""
# The optimizer is created in self.__init__, but the
# hyperparameters need to be set up with self.adjust_optimizer
self.adjust_optimizer(lr=lr,
history_size=history_size,
line_search_fn=line_search_fn)
# Now, we run the optimize routine defined in the base class
return super(LBFGSReconstructor, self).optimize(
iterations,
batch_size=len(self.dataset),
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width)
src/cdtools/reconstructors/sgd.py
+164
View File
@@ -0,0 +1,164 @@
"""This module contains the SGDReconstructor subclass for performing
optimization ('reconstructions') on ptychographic/CDI models using
stochastic gradient descent.
The Reconstructor class is designed to resemble so-called
'Trainer' classes that (in the language of the AI/ML folks) handles
the 'training' of a model given some dataset and optimizer.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch as t
from typing import List, Union
from cdtools.reconstructors import Reconstructor
if TYPE_CHECKING:
from cdtools.models import CDIModel
from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
__all__ = ['SGDReconstructor']
class SGDReconstructor(Reconstructor):
"""
The SGDReconstructor subclass handles the optimization ('reconstruction')
of ptychographic models and datasets using the SGD optimizer.
Parameters
----------
model: CDIModel
Model for CDI/ptychography reconstruction.
dataset: Ptycho2DDataset
The dataset to reconstruct against.
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use.
Important attributes:
- **model** -- Always points to the core model used.
- **optimizer** -- This class by default uses `torch.optim.Adam` to perform
optimizations.
- **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the
`optimize` method.
- **data_loader** -- A torch.utils.data.DataLoader that is defined by
calling the `setup_dataloader` method.
"""
def __init__(self,
model: CDIModel,
dataset: Ptycho2DDataset,
subset: List[int] = None):
# Define the optimizer for use in this subclass
optimizer = t.optim.SGD(model.parameters())
super().__init__(
model,
dataset,
optimizer,
subset=subset,
)
def adjust_optimizer(self,
lr: int = 0.005,
momentum: float = 0,
dampening: float = 0,
weight_decay: float = 0,
nesterov: bool = False):
"""
Change hyperparameters for the utilized optimizer.
Parameters
----------
lr : float
Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
is typically the highest possible value with any chance of being
stable.
momentum : float
Optional, the length of the history to use.
dampening : float
Optional, dampening for the momentum.
weight_decay : float
Optional, weight decay (L2 penalty).
nesterov : bool
Optional, enables Nesterov momentum. Only applicable when momentum
is non-zero.
"""
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
param_group['momentum'] = momentum
param_group['dampening'] = dampening
param_group['weight_decay'] = weight_decay
param_group['nesterov'] = nesterov
def optimize(self,
iterations: int,
batch_size: int = 15,
lr: float = 2e-7,
momentum: float = 0,
dampening: float = 0,
weight_decay: float = 0,
nesterov: bool = False,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10,
shuffle: bool = True):
"""
Runs a round of reconstruction using the Adam optimizer
Formerly `CDIModel.Adam_optimize`
This calls the Reconstructor.optimize superclass method
(formerly `CDIModel.AD_optimize`) to run a round of reconstruction
once the dataloader and optimizer hyperparameters have been
set up.
Parameters
----------
iterations : int
How many epochs of the algorithm to run.
batch_size : int
Optional, the size of the minibatches to use.
lr : float
Optional, The learning rate to use. The default is 2e-7.
momentum : float
Optional, the length of the history to use.
dampening : float
Optional, dampening for the momentum.
weight_decay : float
Optional, weight decay (L2 penalty).
nesterov : bool
Optional, enables Nesterov momentum. Only applicable when momentum
is non-zero.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. Does not affect the result, only
the calculation speed.
shuffle : bool
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
"""
# The optimizer is created in self.__init__, but the
# hyperparameters need to be set up with self.adjust_optimizer
self.adjust_optimizer(lr=lr,
momentum=momentum,
dampening=dampening,
weight_decay=weight_decay,
nesterov=nesterov)
# Now, we run the optimize routine defined in the base class
return super(SGDReconstructor, self).optimize(
iterations,
batch_size=batch_size,
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width,
)
src/cdtools/tools/analysis/analysis.py
+19 -5
View File
@@ -14,6 +14,7 @@ from scipy import linalg as sla
from scipy import special
from scipy import optimize as opt
from scipy import spatial
import warnings
__all__ = [
'product_svd',
@@ -1443,6 +1444,13 @@ def calc_spectral_info(dataset, nbins=50):
the scan pattern whose area matches one detector conjugate field of
view.
This estimation will start to deviate from the truth if the scan area
is not significantly larger than the illumination function, because
the nonzero size of the illumination function is not taken into account.
Furthermore, in the edge case where all the scan points are colinear,
the estimate will fail, and the mean diffraction pattern will be returned
instead
Parameters
----------
dataset : Ptycho2DDataset
@@ -1461,10 +1469,12 @@ def calc_spectral_info(dataset, nbins=50):
"""
scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
scan_area = scan_hull.volume
try:
scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
scan_area = scan_hull.volume
except spatial._qhull.QhullError as e:
scan_area = None
ewg = cdtools.tools.initializers.exit_wave_geometry
obj_basis = ewg(
dataset.detector_geometry['basis'],
@@ -1477,8 +1487,12 @@ def calc_spectral_info(dataset, nbins=50):
np.cross(obj_basis[:,0]*dataset.patterns.shape[-2],
obj_basis[:,1]*dataset.patterns.shape[-1])
)
scale_factor = det_conj_fov_area / scan_area
if scan_area is not None:
scale_factor = det_conj_fov_area / scan_area
else:
warnings.warn("The scan points in this dataset are all colinear. The mean pattern will be calculated rather than a scaled mean based on the scanned area.")
scale_factor = 1/len(dataset)
mask = dataset.mask.cpu().numpy().astype(int)
sum_pattern = dataset.mask * t.sum(dataset.patterns, dim=0) * scale_factor
src/cdtools/tools/data/data.py
+75
View File
@@ -22,6 +22,7 @@ __all__ = ['get_entry_info',
'get_wavelength',
'get_detector_geometry',
'get_mask',
'get_qe_mask',
'get_dark',
'get_data',
'get_shot_to_shot_info',
@@ -32,6 +33,7 @@ __all__ = ['get_entry_info',
'add_source',
'add_detector',
'add_mask',
'add_qe_mask',
'add_dark',
'add_data',
'add_shot_to_shot_info',
@@ -300,6 +302,42 @@ def get_mask(cxi_file):
return None
def get_qe_mask(cxi_file):
"""Returns the quantum efficiency mask defined in the cxi file object
There is no way to store a quantum efficiency mask (a.k.a. a flat-field
image) in the .cxi file specification, but experience has indicated that
this is often a valuable thing to store, because just correcting for a
flatfield with e.g. a division will mess up the photon counting statistics.
Because there is no specification, I have simply chosen to store the
quantum efficiency mask as a float32 array in the same location as the
mask is, i.e. `entry_1/instrument_1/detector_1/qe_mask`.
The stored quantum efficiency mask should be defined as the mask that
a simulated intensity pattern needs to be multiplied by to realize the
measured image. In other words, it should be a flat-field image, not the
inverse of a flat-field image.
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns
-------
qe_mask : np.array
A float32 array storing the quantum efficiency mask from the cxi file
"""
i1 = cxi_file['entry_1/instrument_1']
if 'detector_1/qe_mask' in i1:
qe_mask = i1['detector_1/qe_mask'][()].astype(np.float32)
return qe_mask
else:
return None
def get_dark(cxi_file):
"""Returns an array with a dark image to use for initialization of a background model
@@ -635,6 +673,43 @@ def add_mask(cxi_file, mask):
d1.create_dataset('mask',data=mask_to_save)
def add_qe_mask(cxi_file, qe_mask):
"""Adds the specified quantum efficiency mask to the cxi file
There is no way to store a quantum efficiency mask (a.k.a. a flat-field
image) in the .cxi file specification, but experience has indicated that
this is often a valuable thing to store, because just correcting for a
flatfield with e.g. a division will mess up the photon counting statistics.
Because there is no specification, I have simply chosen to store the
quantum efficiency mask as an array in the same location as the
mask is, i.e. `entry_1/instrument_1/detector_1/qe_mask`.
The stored quantum efficiency mask should be defined as the mask that
a simulated intensity pattern needs to be multiplied by to realize the
measured image. In other words, it should be a flat-field image, not the
inverse of a flat-field image.
Parameters
----------
cxi_file : h5py.File
The file to add the mask to
qe_mask : array
The quantum efficiency mask to save out to the file
"""
if 'entry_1/instrument_1' not in cxi_file:
cxi_file['entry_1'].create_group('instrument_1')
i1 = cxi_file['entry_1/instrument_1']
if 'detector_1' not in i1:
i1.create_group('detector_1')
d1 = i1['detector_1']
if isinstance(qe_mask, t.Tensor):
qe_mask = qe_mask.detach().cpu().numpy()
d1.create_dataset('qe_mask',data=qe_mask)
def add_dark(cxi_file, dark):
"""Adds the specified dark image to a cxi file
src/cdtools/tools/measurements/measurements.py
+27 -11
View File
@@ -155,7 +155,18 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
return t.clamp(output + epsilon,0,saturation)
def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1, simulate_finite_pixels=False):
def quadratic_background(
wavefield,
background,
*args,
detector_slice=None,
measurement=intensity,
epsilon=1e-7,
qe_mask=None,
saturation=None,
oversampling=1,
simulate_finite_pixels=False
):
"""Returns the intensity of a wavefield plus a background
The intensity is calculated via the given measurment function
@@ -173,6 +184,8 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
Optional, a slice or tuple of slices defining a section of the simulation to return
measurement : function
Default is measurements.intensity, the measurement function to use.
qe_mask : torch.Tensor
A tensor storing the per-pixel quantum efficiency (up to an unknown global scaling factor)
saturation : float
Optional, a maximum saturation value to clamp the resulting intensities to
oversampling : int
@@ -184,17 +197,20 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
A real MxN array storing the wavefield's intensities
"""
if detector_slice is None:
output = measurement(wavefield, *args, epsilon=epsilon,
oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels) \
+ background**2
else:
output = measurement(wavefield, *args, detector_slice=detector_slice,
epsilon=epsilon, oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels) \
+ background**2
raw_intensity = measurement(
wavefield,
*args,
detector_slice=detector_slice,
epsilon=epsilon,
oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels
)
if qe_mask is None:
output = raw_intensity + background**2
else:
output = (qe_mask * raw_intensity) + background**2
# This has to be done after the background is added, hence we replicate
# it here
if saturation is None:
src/cdtools/tools/plotting/plotting.py
+24 -6
View File
@@ -522,7 +522,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
units=units, show_cbar=False, **kwargs)
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs):
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, clear_fig=True, label=None, color=None, marker='.', **kwargs):
"""Plots a set of probe translations in a nicely formatted way
Parameters
@@ -537,6 +537,14 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
Whether to plot lines indicating the path taken
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
clear_fig : bool
Default is True. Whether to clear the figure before plotting.
label : str
Default is None. A label to give the plotted markers for a legend.
color : str
Default is None. The color to plot the markers in. By default, will follow the matplotlib color cycle.
color : str
Default is '.'. The marker style to plot with.
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -554,18 +562,28 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if clear_fig:
plt.gcf().clear()
if isinstance(translations, t.Tensor):
translations = translations.detach().cpu().numpy()
translations = translations * factor
plt.plot(translations[:,0], translations[:,1],'k.')
linestyle = '-' if lines else 'None'
linewidth = 1 if lines else 0
plt.plot(translations[:,0], translations[:,1],
marker=marker, linestyle=linestyle,
label=label, color=color,
linewidth=linewidth)
if invert_xaxis:
plt.gca().invert_xaxis()
ax = plt.gca()
x_min, x_max = ax.get_xlim()
# Protect against flipping twice if plotting on top of existing graph
if x_min <= x_max:
ax.invert_xaxis()
if lines:
plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
tests/conftest.py
+87 -61
View File
@@ -1,9 +1,9 @@
import numpy as np
import torch as t
import h5py
import pytest
import datetime
import h5py
import numpy as np
import pytest
import torch as t
#
#
@@ -47,6 +47,7 @@ def pytest_collection_modifyitems(config, items):
if "slow" in item.keywords:
item.add_marker(skip_slow)
@pytest.fixture
def reconstruction_device(request):
return request.config.getoption("--reconstruction_device")
@@ -57,7 +58,6 @@ def show_plot(request):
return request.config.getoption("--plot")
@pytest.fixture(scope='module')
def ptycho_cxi_1():
"""Creates an example file for CXI ptychography. This file is defined
@@ -67,11 +67,11 @@ def ptycho_cxi_1():
"""
expected = {}
f = h5py.File('ptycho_cxi_1','w',driver='core',backing_store=False)
f = h5py.File('ptycho_cxi_1', 'w', driver='core', backing_store=False)
# Start by defining the basic structure
f.create_dataset('cxi_version', data=150)
f.create_dataset('number_of_entries',data=1)
f.create_dataset('number_of_entries', data=1)
# Then define a bunch of metadata for entry_1
e1f = f.create_group('entry_1')
@@ -104,19 +104,19 @@ def ptycho_cxi_1():
s1f['concentration'] = s1e['concentration']
s1e['mass'] = np.float32(np.random.rand())
s1f['mass'] = s1e['mass']
s1e['temperature'] = np.float32(np.random.rand()*100)
s1e['temperature'] = np.float32(np.random.rand() * 100)
s1f['temperature'] = s1e['temperature']
s1e['thickness'] = np.float32(np.random.rand()*1e-7)
s1e['thickness'] = np.float32(np.random.rand() * 1e-7)
s1f['thickness'] = s1e['thickness']
s1e['unit_cell_volume'] = np.float32(np.random.rand() * 1e-27)
s1f['unit_cell_volume'] = s1e['unit_cell_volume']
s1e['unit_cell'] = np.array([1,1,1,90,90,90]).astype(np.float32)
s1f.create_dataset('unit_cell',data = s1e['unit_cell'])
s1e['unit_cell'] = np.array([1, 1, 1, 90, 90, 90]).astype(np.float32)
s1f.create_dataset('unit_cell', data=s1e['unit_cell'])
i1f = e1f.create_group('instrument_1')
source1f = i1f.create_group('source_1')
energy = np.float32(1.3618e-16) #Joules, = 850 eV
energy = np.float32(1.3618e-16) # Joules, = 850 eV
source1f['energy'] = energy
expected['wavelength'] = np.float32(1.9864459e-25) / energy
source1f['wavelength'] = expected['wavelength']
@@ -126,41 +126,48 @@ def ptycho_cxi_1():
d1e = expected['detector']
d1e['distance'] = np.float32(0.3)
d1f['distance'] = d1e['distance']
d1e['basis'] = np.array([[0,-30e-6,0],
[-20e-6,0,0]]).astype(np.float32).transpose()
d1f.create_dataset('basis_vectors',data=d1e['basis'])
d1e['basis'] = np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).astype(np.float32).transpose()
d1f.create_dataset('basis_vectors', data=d1e['basis'])
d1f['x_pixel_size'] = np.float32(20e-6)
d1f['y_pixel_size'] = np.float32(30e-6)
d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
d1e['corner'] = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
d1f.create_dataset('corner_position', data=d1e['corner'])
# Remember the format for the CXI file differs from the format used
# internally
mask = np.zeros((256,256)).astype(np.int32)
expected['mask'] = np.ones((256,256)).astype(bool)
d1f.create_dataset('mask',data=mask)
mask = np.zeros((256, 256)).astype(np.int32)
mask[5, 8] = 1
expected['mask'] = np.ones((256, 256)).astype(bool)
expected['mask'][5, 8] = 0
d1f.create_dataset('mask', data=mask)
# There is no specification for this in the CXI file format :(
qe_mask = np.ones((256, 256)).astype(np.float32)
expected['qe_mask'] = qe_mask
d1f.create_dataset('qe_mask', data=qe_mask)
# Create an initial background
dark = np.ones((256,256)) * 0.01
dark = np.ones((256, 256)) * 0.01
expected['dark'] = dark
d1f.create_dataset('data_dark', data=dark)
data1f = e1f.create_group('data_1')
data = np.random.rand(100,256,256).astype(np.float32)
data = np.random.rand(100, 256, 256).astype(np.float32)
expected['data'] = data
d1f.create_dataset('data',data=data)
d1f.create_dataset('data', data=data)
data1f['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data')
d1f['data'].attrs['axes'] = np.bytes_('translation:y:x')
expected['axes'] = ['translation','y','x']
expected['axes'] = ['translation', 'y', 'x']
g1f = s1f.create_group('geometry_1')
orientation = np.array([1.,0,0,0,1,0])
orientation = np.array([1., 0, 0, 0, 1, 0])
g1f.create_dataset('orientation', data=orientation)
s1e['orientation'] = np.array([[1.,0,0],[0,1,0],[0,0,1]])
translations = np.arange(300).reshape((100,3)).astype(np.float32)
g1f.create_dataset('translation',data=translations)
s1e['orientation'] = np.array([[1., 0, 0], [0, 1, 0], [0, 0, 1]])
translations = np.arange(300).reshape((100, 3)).astype(np.float32)
g1f.create_dataset('translation', data=translations)
data1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
d1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
expected['translations'] = -translations
@@ -186,11 +193,11 @@ def ptycho_cxi_2():
"""
expected = {}
f = h5py.File('ptycho_cxi_2','w',driver='core',backing_store=False)
f = h5py.File('ptycho_cxi_2', 'w', driver='core', backing_store=False)
# Start by defining the basic structure
f.create_dataset('cxi_version', data=150)
f.create_dataset('number_of_entries',data=1)
f.create_dataset('number_of_entries', data=1)
# Then define a bunch of metadata for entry_1
e1f = f.create_group('entry_1')
@@ -203,13 +210,13 @@ def ptycho_cxi_2():
s1f = e1f.create_group('sample_1')
expected['sample info'] = {}
s1e = expected['sample info']
s1e['temperature'] = np.float32(np.random.rand()*100)
s1e['temperature'] = np.float32(np.random.rand() * 100)
s1f['temperature'] = s1e['temperature']
i1f = e1f.create_group('instrument_1')
source1f = i1f.create_group('source_1')
energy = np.float32(1.3618e-16) #Joules, = 850 eV
energy = np.float32(1.3618e-16) # Joules, = 850 eV
expected['wavelength'] = np.float32(1.9864459e-25) / energy
source1f['wavelength'] = expected['wavelength']
@@ -217,34 +224,35 @@ def ptycho_cxi_2():
expected['detector'] = {}
d1e = expected['detector']
d1e['distance'] = np.float32(0.3)
d1e['basis'] = np.array([[0,-30e-6,0],
[-20e-6,0,0]]).astype(np.float32).transpose()
d1e['basis'] = np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).astype(np.float32).transpose()
d1f['x_pixel_size'] = np.float32(20e-6)
d1f['y_pixel_size'] = np.float32(30e-6)
d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
d1e['corner'] = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
d1f.create_dataset('corner_position', data=d1e['corner'])
# Remember the format for the CXI file differs from the format used
# internally
expected['mask'] = None
expected['qe_mask'] = None
# Test with a set of dark images
dark = np.ones((10,256,256)) * 0.01
expected['dark'] = np.nanmean(dark,axis=0)
dark = np.ones((10, 256, 256)) * 0.01
expected['dark'] = np.nanmean(dark, axis=0)
d1f.create_dataset('data_dark', data=dark)
e1f.create_group('data_1')
data1f = e1f.create_group('data_1')
data = np.random.rand(100,256,256).astype(np.float32)
data = np.random.rand(100, 256, 256).astype(np.float32)
expected['data'] = data
d1f.create_dataset('data',data=data)
d1f.create_dataset('data', data=data)
expected['axes'] = None
g1f = s1f.create_group('geometry_1')
translations = np.arange(300).reshape((100,3)).astype(np.float32)
g1f.create_dataset('translation',data=translations)
translations = np.arange(300).reshape((100, 3)).astype(np.float32)
g1f.create_dataset('translation', data=translations)
expected['translations'] = -translations
yield f, expected
@@ -267,11 +275,11 @@ def ptycho_cxi_3():
"""
expected = {}
f = h5py.File('ptycho_cxi_3','w',driver='core',backing_store=False)
f = h5py.File('ptycho_cxi_3', 'w', driver='core', backing_store=False)
# Start by defining the basic structure
f.create_dataset('cxi_version', data=150)
f.create_dataset('number_of_entries',data=1)
f.create_dataset('number_of_entries', data=1)
# Then define a bunch of metadata for entry_1
e1f = f.create_group('entry_1')
@@ -288,7 +296,7 @@ def ptycho_cxi_3():
i1f = e1f.create_group('instrument_1')
source1f = i1f.create_group('source_1')
energy = np.float32(1.3618e-16) #Joules, = 850 eV
energy = np.float32(1.3618e-16) # Joules, = 850 eV
source1f['energy'] = energy
expected['wavelength'] = np.float32(1.9864459e-25) / energy
@@ -297,36 +305,41 @@ def ptycho_cxi_3():
d1e = expected['detector']
d1e['distance'] = np.float32(0.3)
d1f['distance'] = d1e['distance']
d1e['basis'] = np.array([[0,-30e-6,0],
[-20e-6,0,0]]).astype(np.float32).transpose()
d1f.create_dataset('basis_vectors',data=d1e['basis'])
d1e['basis'] = np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).astype(np.float32).transpose()
d1f.create_dataset('basis_vectors', data=d1e['basis'])
d1e['corner'] = None
# Remember the format for the CXI file differs from the format used
# internally
mask = np.ones((256,256)).astype(np.uint32) * 0x00001000
expected['mask'] = np.ones((256,256)).astype(bool)
d1f.create_dataset('mask',data=mask)
mask = np.ones((256, 256)).astype(np.uint32) * 0x00001000
mask[15, 47] = 38
expected['mask'] = np.ones((256, 256)).astype(bool)
expected['mask'][15, 47] = 0
d1f.create_dataset('mask', data=mask)
expected['qe_mask'] = None
expected['dark'] = None
data1f = e1f.create_group('data_1')
data = np.random.rand(100,256,256).astype(np.float32)
data = np.random.rand(100, 256, 256).astype(np.float32)
expected['data'] = data
data1f.create_dataset('data',data=data)
data1f.create_dataset('data', data=data)
data1f['data'].attrs['axes'] = np.bytes_('translation:y:x')
expected['axes'] = ['translation','y','x']
expected['axes'] = ['translation', 'y', 'x']
translations = np.arange(300).reshape((100,3)).astype(np.float32)
data1f.create_dataset('translation',data=translations)
translations = np.arange(300).reshape((100, 3)).astype(np.float32)
data1f.create_dataset('translation', data=translations)
expected['translations'] = -translations
yield f, expected
f.close()
@pytest.fixture(scope='module')
def polarized_ptycho_cxi(ptycho_cxi_1):
f, expected = ptycho_cxi_1
@@ -338,7 +351,7 @@ def polarized_ptycho_cxi(ptycho_cxi_1):
data1f.create_dataset('polarizer_angle', data=expected['polarizer_angle'])
yield f, expected
# As specific issues start to crop up with loading CXI files from different
# beamlines, put a fixture here that replicates the issue so that we can
@@ -360,16 +373,29 @@ def gold_ball_cxi(pytestconfig):
return str(pytestconfig.rootpath) + \
'/examples/example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
@pytest.fixture(scope='module')
def lab_ptycho_cxi(pytestconfig):
return str(pytestconfig.rootpath) + \
'/examples/example_data/lab_ptycho_data.cxi'
@pytest.fixture(scope='module')
def optical_data_ss_cxi(pytestconfig):
return str(pytestconfig.rootpath) + \
'/examples/example_data/Optical_Data_ss.cxi'
@pytest.fixture(scope='module')
def optical_ptycho_incoherent_pickle(pytestconfig):
return str(pytestconfig.rootpath) + \
'/examples/example_data/Optical_ptycho_incoherent.pickle'
@pytest.fixture(scope='module')
def example_nested_dicts(pytestconfig):
example_tensor = t.as_tensor(np.array([1,4.5,7]))
example_array = np.ones([10,20,30])
example_tensor = t.as_tensor(np.array([1, 4.5, 7]))
example_array = np.ones([10, 20, 30])
example_scalar = 4.5
example_single_element_array = np.array([0.3])
example_string = 'testing'
tests/models/test_fancy_ptycho.py
+60 -68
View File
@@ -1,28 +1,74 @@
import pytest
import cdtools
import torch as t
import cdtools
from matplotlib import pyplot as plt
# Force all reconstructions to use the same RNG seed
t.manual_seed(0)
def test_center_probe(lab_ptycho_cxi):
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
fourier_probe=False
)
base_probe = model.probe.detach().clone()
model.center_probes()
centered_probe = model.probe.detach().clone()
fourier_model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
fourier_probe=True,
)
fourier_model.probe.data = cdtools.tools.propagators.far_field(
base_probe
)
fourier_model.probe.detach().clone()
fourier_model.center_probes()
fourier_centered_probe = fourier_model.probe.detach().clone()
ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
fourier_centered_probe)
# So we know the code had to do something
assert not t.allclose(base_probe, centered_probe)
# And checking that they both do the same thing, whether or not
# fourier_probe was set to True
assert t.allclose(
centered_probe,
ifft_fourier_centered_probe,
atol=1e-4,
rtol=1e-3
)
@pytest.mark.slow
def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
print('\nTesting performance on the standard transmission ptycho dataset')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
# Test the masking system
dataset.mask[110:115,65:70] = 0
dataset.patterns[...,~dataset.mask] = t.max(dataset.patterns)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
n_modes=3,
oversampling=2,
exponentiate_obj=True,
dm_rank=2,
exponentiate_obj=True,
probe_support_radius=120,
propagation_distance=5e-3,
units='mm',
propagation_distance=5e-3,
units='mm',
obj_view_crop=-50,
use_qe_mask=True, # test this in the case where no qe mask is defined
)
print('Running reconstruction on provided reconstruction_device,',
reconstruction_device)
model.to(device=reconstruction_device)
@@ -33,11 +79,16 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50):
for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(25, dataset, lr=0.001, batch_size=50):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
model.tidy_probes()
if show_plot:
@@ -45,64 +96,5 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
model.compare(dataset)
# If this fails, the reconstruction has gotten worse
assert model.loss_history[-1] < 0.001
@pytest.mark.slow
def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
print('\nTesting performance on the standard gold balls dataset')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
pad = 10
dataset.pad(pad)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
probe_support_radius=50,
propagation_distance=2e-6,
units='um',
probe_fourier_crop=pad
)
model.translation_offsets.data += \
0.7 * t.randn_like(model.translation_offsets)
# Not much probe intensity instability in this dataset, no need for this
model.weights.requires_grad = False
print('Running reconstruction on provided --reconstruction_device,',
reconstruction_device)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100,
schedule=True):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# This just comes from running a reconstruction when it was working well
# and choosing a rough value. If it triggers this assertion error,
# something changed to make the final quality worse!
assert model.loss_history[-1] < 0.0001
assert model.loss_history[-1] < 0.0013
tests/models/test_simple_ptycho.py
+9 -4
View File
@@ -1,13 +1,18 @@
import pytest
import torch as t
import cdtools
from matplotlib import pyplot as plt
# Force all reconstructions to use the same RNG seed
t.manual_seed(0)
@pytest.mark.slow
def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
model = cdtools.models.SimplePtycho.from_dataset(dataset)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
@@ -15,10 +20,10 @@ def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# If this fails, the reconstruction got worse
assert model.loss_history[-1] < 0.013
tests/test_datasets.py
+178 -102
View File
@@ -1,12 +1,16 @@
import datetime
import itertools
import os
from copy import deepcopy
import h5py
import numpy as np
import pytest
import torch as t
from cdtools.datasets import CDataset, Ptycho2DDataset
from cdtools.tools import data as cdtdata
import numpy as np
import torch as t
import h5py
import datetime
from copy import deepcopy
import pytest
import itertools
#
# We start by testing the CDataset base class
@@ -14,20 +18,20 @@ import itertools
def test_CDataset_init():
entry_info = {'start_time': datetime.datetime.now(),
'title' : 'A simple test'}
'title': 'A simple test'}
sample_info = {'name': 'A test sample',
'mass' : 3.4,
'unit_cell' : np.array([1,1,1,87,84.5,90])}
'mass': 3.4,
'unit_cell': np.array([1, 1, 1, 87, 84.5, 90])}
wavelength = 1e-9
detector_geometry = {'distance': 0.7,
'basis': np.array([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(),
'corner': np.array((2550e-6,3825e-6,0.3))}
mask = np.ones((256,256))
'basis': np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).transpose(),
'corner': np.array((2550e-6, 3825e-6, 0.3))}
mask = np.ones((256, 256))
dataset = CDataset(entry_info, sample_info,
wavelength, detector_geometry, mask)
assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(bool))))
assert t.all(t.eq(dataset.mask, t.tensor(mask.astype(bool))))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
@@ -49,7 +53,7 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
else:
assert dataset.sample_info is not None
assert np.isclose(dataset.wavelength,expected['wavelength'])
assert np.isclose(dataset.wavelength, expected['wavelength'])
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
@@ -60,15 +64,16 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
assert 'corner' in dataset.detector_geometry
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
assert t.all(t.eq(t.tensor(expected['mask']), dataset.mask))
if expected['qe_mask'] is not None:
assert t.all(t.eq(t.tensor(expected['qe_mask']), dataset.qe_mask))
if expected['dark'] is not None:
assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
dataset.background))
def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
for cxi, expected in test_ptycho_cxis:
dataset = CDataset.from_cxi(cxi)
@@ -88,24 +93,24 @@ def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
# Just check one of the loaded attributes
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
read_dataset.detector_geometry['distance'])
# Check that the other ones are loaded but not for fidelity
assert 'basis' in read_dataset.detector_geometry
if dataset.detector_geometry['corner'] is not None:
assert 'corner' in read_dataset.detector_geometry
if dataset.mask is not None:
assert t.all(t.eq(dataset.mask,read_dataset.mask))
assert t.all(t.eq(dataset.mask, read_dataset.mask))
if dataset.qe_mask is not None:
assert t.all(t.eq(dataset.qe_mask, read_dataset.qe_mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
def test_CDataset_to(ptycho_cxi_1):
dataset = CDataset.from_cxi(ptycho_cxi_1[0])
@@ -115,6 +120,7 @@ def test_CDataset_to(ptycho_cxi_1):
if t.cuda.is_available():
dataset.to(device='cuda:0')
assert dataset.mask.device == t.device('cuda:0')
assert dataset.qe_mask.device == t.device('cuda:0')
assert dataset.background.device == t.device('cuda:0')
@@ -125,27 +131,46 @@ def test_CDataset_to(ptycho_cxi_1):
def test_Ptycho2DDataset_init():
entry_info = {'start_time': datetime.datetime.now(),
'title' : 'A simple test'}
'title': 'A simple test'}
sample_info = {'name': 'A test sample',
'mass' : 3.4,
'unit_cell' : np.array([1,1,1,87,84.5,90])}
'mass': 3.4,
'unit_cell': np.array([1, 1, 1, 87, 84.5, 90])}
wavelength = 1e-9
detector_geometry = {'distance': 0.7,
'basis': np.array([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(),
'corner': np.array((2550e-6,3825e-6,0.3))}
mask = np.ones((256,256))
patterns = np.random.rand(20,256,256)
translations = np.random.rand(20,3)
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
'basis': np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).transpose(),
'corner': np.array((2550e-6, 3825e-6, 0.3))}
mask = np.ones((256, 256))
qe_mask = 1.2 * np.ones((256, 256), dtype=np.float32)
patterns = np.random.rand(20, 256, 256)
translations = np.random.rand(20, 3)
assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
assert t.all(t.eq(dataset.mask, t.BoolTensor(mask)))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
assert dataset.detector_geometry == detector_geometry
assert t.allclose(dataset.patterns, t.as_tensor(patterns))
assert t.allclose(dataset.translations, t.as_tensor(translations))
# Also test one with a qe_mask
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask,
qe_mask=qe_mask)
assert t.all(t.eq(dataset.mask, t.BoolTensor(mask)))
assert t.all(t.eq(dataset.qe_mask, t.as_tensor(qe_mask)))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
@@ -169,7 +194,7 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
else:
assert dataset.sample_info is not None
assert np.isclose(dataset.wavelength,expected['wavelength'])
assert np.isclose(dataset.wavelength, expected['wavelength'])
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
@@ -180,15 +205,51 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
assert 'corner' in dataset.detector_geometry
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
assert t.all(t.eq(t.tensor(expected['mask']), dataset.mask))
if expected['qe_mask'] is not None:
assert t.all(t.eq(t.tensor(expected['qe_mask']), dataset.qe_mask))
if expected['dark'] is not None:
assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
dataset.background))
assert t.allclose(t.tensor(expected['data']),dataset.patterns)
assert t.allclose(t.tensor(expected['translations']),dataset.translations)
assert t.allclose(t.tensor(expected['data']), dataset.patterns)
assert t.allclose(t.tensor(expected['translations']), dataset.translations)
def test_Ptycho2DDataset_from_cxi_64bit(test_ptycho_cxis):
"""Test that we can load a 64-bit cxi file. Should issue
a warning, but still load the data."""
# create test patterns and translations
np.random.seed(42)
patterns = np.random.rand(20, 256, 256).astype(np.float64)
translations = np.random.rand(20, 3).astype(np.float64)
dataset = Ptycho2DDataset(translations, patterns)
dataset.detector_geometry = {
'distance': 0.1, # in meters
'basis': t.tensor([
[-0e-06, -13.5e-06 * 4],
[-13.5e-06 * 4, 0e-06],
[0e-06, 0e-06]
]),
'corner': None
}
dataset.wavelength = 1.6891579427792915e-09 # in meters
# and save to a temp file
dataset.to_cxi('test_Ptycho2DDataset_from_cxi_64bit.cxi')
with pytest.warns(UserWarning, match='64-bit floats'):
dataset_64bit = Ptycho2DDataset.from_cxi('test_Ptycho2DDataset_from_cxi_64bit.cxi')
# Check that the data is loaded correctly
assert dataset_64bit.patterns.dtype == t.float32
assert dataset_64bit.translations.dtype == t.float32
# delete the created test file
os.remove('test_Ptycho2DDataset_from_cxi_64bit.cxi')
def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
@@ -212,19 +273,19 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
# Just check one of the loaded attributes
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
read_dataset.detector_geometry['distance'])
# Check that the other ones are loaded but not for fidelity
assert 'basis' in read_dataset.detector_geometry
if dataset.detector_geometry['corner'] is not None:
assert 'corner' in read_dataset.detector_geometry
if dataset.mask is not None:
assert t.all(t.eq(dataset.mask,read_dataset.mask))
assert t.all(t.eq(dataset.mask, read_dataset.mask))
if dataset.qe_mask is not None:
assert t.all(t.eq(dataset.qe_mask, read_dataset.qe_mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
@@ -235,15 +296,16 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
def test_Ptycho2DDataset_to(ptycho_cxi_1):
dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
dataset.to(dtype=t.float64)
assert dataset.mask.dtype == t.bool
assert dataset.qe_mask.dtype == t.float64
assert dataset.patterns.dtype == t.float64
assert dataset.translations.dtype == t.float64
# If cuda is available, check that moving the mask to CUDA works.
if t.cuda.is_available():
dataset.to(device='cuda:0')
assert dataset.mask.device == t.device('cuda:0')
assert dataset.qe_mask.device == t.device('cuda:0')
assert dataset.background.device == t.device('cuda:0')
assert dataset.patterns.device == t.device('cuda:0')
assert dataset.translations.device == t.device('cuda:0')
@@ -257,8 +319,8 @@ def test_Ptycho2DDataset_ops(ptycho_cxi_1):
assert len(dataset) == expected['data'].shape[0]
(idx, translation), pattern = dataset[3]
assert idx == 3
assert t.allclose(translation, t.tensor(expected['translations'][3,:]))
assert t.allclose(pattern, t.tensor(expected['data'][3,:,:]))
assert t.allclose(translation, t.tensor(expected['translations'][3, :]))
assert t.allclose(pattern, t.tensor(expected['data'][3, :, :]))
def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
@@ -271,12 +333,12 @@ def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
(idx, translation), pattern = dataset[3]
assert str(translation.device) == 'cuda:0'
assert str(pattern.device) == 'cuda:0'
assert idx == 3
assert t.allclose(translation.to(device='cpu'),
t.tensor(expected['translations'][3,:]))
t.tensor(expected['translations'][3, :]))
assert t.allclose(pattern.to(device='cpu'),
t.tensor(expected['data'][3,:,:]))
t.tensor(expected['data'][3, :, :]))
def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
@@ -291,34 +353,51 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
# May start failing if the test datasets are changed to include
# a dataset with any dimension not even. That's a problem with the
# test, not the code. Sorry! -Abe
masked_patterns = dataset.mask * dataset.patterns
assert t.allclose(
copied_dataset.patterns,
dataset.patterns[:,::2,::2] +
dataset.patterns[:,1::2,::2] +
dataset.patterns[:,::2,1::2] +
dataset.patterns[:,1::2,1::2]
masked_patterns[:, ::2, ::2] + masked_patterns[:, 1::2, ::2] + masked_patterns[:, ::2, 1::2] + masked_patterns[:, 1::2, 1::2]
)
assert t.allclose(
copied_dataset.mask,
t.logical_and(
t.logical_and(dataset.mask[::2,::2],
dataset.mask[1::2,::2]),
t.logical_and(dataset.mask[::2,1::2],
dataset.mask[1::2,1::2]),
if dataset.qe_mask is None:
manually_downsampled_mask = t.logical_and(
t.logical_and(dataset.mask[::2, ::2],
dataset.mask[1::2, ::2]),
t.logical_and(dataset.mask[::2, 1::2],
dataset.mask[1::2, 1::2])
)
assert t.allclose(
copied_dataset.mask,
manually_downsampled_mask,
)
else:
manually_downsampled_mask = t.logical_or(
t.logical_or(dataset.mask[::2, ::2],
dataset.mask[1::2, ::2]),
t.logical_or(dataset.mask[::2, 1::2],
dataset.mask[1::2, 1::2])
)
assert t.allclose(
copied_dataset.mask,
manually_downsampled_mask
)
)
masked_qe_mask = dataset.mask * dataset.qe_mask
manually_downsampled_qe_mask = (
masked_qe_mask[::2, ::2] + masked_qe_mask[1::2, ::2] + masked_qe_mask[::2, 1::2] + masked_qe_mask[1::2, 1::2]
) / 4
assert t.allclose(
copied_dataset.qe_mask,
manually_downsampled_qe_mask
)
if dataset.background is not None:
assert t.allclose(
copied_dataset.background,
dataset.background[::2,::2] +
dataset.background[1::2,::2] +
dataset.background[::2,1::2] +
dataset.background[1::2,1::2]
)
dataset.background[::2, ::2] + dataset.background[1::2, ::2] + dataset.background[::2, 1::2] + dataset.background[1::2, 1::2]
)
# And then we just test the shape for a few factors, and check that
# it doesn't fail on edge cases (e.g. factor=1)
@@ -333,10 +412,10 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
assert np.allclose(expected_pattern_shape,
np.array(copied_dataset.patterns.shape))
assert np.allclose(np.array(dataset.mask.shape) // factor,
np.array(copied_dataset.mask.shape))
if dataset.background is not None:
assert np.allclose(np.array(dataset.background.shape) // factor,
np.array(copied_dataset.background.shape))
@@ -373,13 +452,13 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
copied_dataset = deepcopy(dataset)
# Test 1: Complain when the the bounds of an ROI are correctly defined,
# but it does not contain any sample positions inside of it. The
# translations in ptycho_cxi_1 ranges from 0m to -300m in both x and y
# (it looks like a line scan). We select an ROI at (0, -300) which should
# but it does not contain any sample positions inside of it. The
# translations in ptycho_cxi_1 ranges from 0m to -300m in both x and y
# (it looks like a line scan). We select an ROI at (0, -300) which should
# not contain any translation positions.
with pytest.raises(ValueError) as excinfo:
copied_dataset.crop_translations(roi=(0,-5,-295,-300))
assert('(i.e., patterns and translations will be empty)') in str(excinfo.value)
copied_dataset.crop_translations(roi=(0, -5, -295, -300))
assert '(i.e., patterns and translations will be empty)' in str(excinfo.value)
# Test 2: Draw an ROI that's centered in the middle of the x/y translation range
# and make sure that the first and last x/y elements in dataset.translate
@@ -387,40 +466,37 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
# Make tuples that will store the x and y positions. This will be used for
# making permutations of the x/y positions in the ROI later.
x_left, y_top = dataset.translations[10,:2]
x_right, y_bottom = dataset.translations[-11,:2]
x_left, y_top = dataset.translations[10, :2]
x_right, y_bottom = dataset.translations[-11, :2]
x_permutations = ((x_left, x_right), (x_right, x_left))
y_permutations = ((y_top, y_bottom), (y_bottom, y_top))
roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in
itertools.product(x_permutations, y_permutations))
roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in itertools.product(x_permutations, y_permutations))
# Get the dataset
copied_dataset.crop_translations(roi=roi_permutations[0])
# Execute the actual test
assert (copied_dataset.translations[0,0] in x_permutations[0]) and \
(copied_dataset.translations[-1,0] in x_permutations[0])
assert (copied_dataset.translations[0,1] in y_permutations[0]) and \
(copied_dataset.translations[-1,1] in y_permutations[0])
assert (copied_dataset.translations[0, 0] in x_permutations[0]) and \
(copied_dataset.translations[-1, 0] in x_permutations[0])
assert (copied_dataset.translations[0, 1] in y_permutations[0]) and \
(copied_dataset.translations[-1, 1] in y_permutations[0])
# Test 3: Check if the shape of dataset.patterns and dataset.translate is correct
# (designed to be 20 fewer rows here)
# In the future, this should include a check for dataset.intensities once
# an appropriate cxi file is set up for conftest.
expected_patterns_shape = np.concatenate([[dataset.patterns.shape[0] - 20],
dataset.patterns.shape[-2:]])
expected_translations_shape = np.concatenate([[dataset.translations.shape[0] - 20],
dataset.translations.shape[-1:]])
expected_patterns_shape = np.concatenate([[dataset.patterns.shape[0] - 20], dataset.patterns.shape[-2:]])
expected_translations_shape = np.concatenate([[dataset.translations.shape[0] - 20], dataset.translations.shape[-1:]])
assert np.allclose(np.array(copied_dataset.patterns.shape), expected_patterns_shape)
assert np.allclose(np.array(copied_dataset.translations.shape), expected_translations_shape)
# Test 4: Make sure that we always get the same result no matter what order we
# define the left/right and bottom/top values in roi, provided that roi[:2]
# define the left/right and bottom/top values in roi, provided that roi[:2]
# and roi[2:] correspond with the x and y coordinates, respectively.
# Check each permutation
@@ -431,6 +507,6 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
copied_dataset.crop_translations(roi=roi)
# Check if the contents of dataset.patterns and dataset.translate is correct
assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10,:])
assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10, :])
assert t.allclose(copied_dataset.translations, dataset.translations[10:-10,:])
assert t.allclose(copied_dataset.translations, dataset.translations[10:-10, :])
tests/test_reconstructors.py
+319
View File
@@ -0,0 +1,319 @@
import pytest
import cdtools
import torch as t
import numpy as np
import pickle
from matplotlib import pyplot as plt
from copy import deepcopy
@pytest.mark.slow
def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
"""
This test checks out several things with the Au particle dataset
1) Calls to Reconstructor.adjust_optimizer is updating the
hyperparameters
2) We are only using the single-GPU dataloading method
3) Ensure `recon.model` points to the original `model`
4) Reconstructions performed by `Adam.optimize` and
`model.Adam_optimize` calls produce identical results when
run over one round of optimization.
5) The quality of the reconstruction remains below a specified
threshold.
5) Ensure that the FancyPtycho model works fine and dandy with the
Reconstructors.
"""
print('\nTesting performance on the standard gold balls dataset ' +
'with reconstructors.AdamReconstructor')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
pad = 10
dataset.pad(pad)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
probe_support_radius=50,
propagation_distance=2e-6,
units='um',
probe_fourier_crop=pad
)
model.translation_offsets.data += 0.7 * \
t.randn_like(model.translation_offsets)
model.weights.requires_grad = False
# Make a copy of the model
model_recon = deepcopy(model)
model.to(device=reconstruction_device)
model_recon.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
# ******* Reconstructions with AdamReconstructor.optimize *******
print('Running reconstruction using AdamReconstructor.optimize' +
' on provided reconstruction_device,', reconstruction_device)
recon = cdtools.reconstructors.AdamReconstructor(model=model_recon,
dataset=dataset)
t.manual_seed(0)
# Run a reconstruction
epoch_tup = (20, 50, 100)
lr_tup = (0.005, 0.002, 0.001)
batch_size_tup = (50, 100, 100)
for i, iterations in enumerate(epoch_tup):
for loss in recon.optimize(iterations,
lr=lr_tup[i],
batch_size=batch_size_tup[i]):
print(model_recon.report())
if show_plot and model_recon.epoch % 10 == 0:
model_recon.inspect(dataset)
# Check hyperparameter update
assert recon.optimizer.param_groups[0]['lr'] == lr_tup[i]
assert recon.data_loader.batch_size == batch_size_tup[i]
# Ensure that recon does not have sampler as an attribute (only used in
# multi-GPU)
assert not hasattr(recon, 'sampler')
# Ensure recon.model points to the original model
assert id(model_recon) == id(recon.model)
model_recon.tidy_probes()
if show_plot:
model_recon.inspect(dataset)
model_recon.compare(dataset)
# ******* Reconstructions with CDIModel.Adam_optimize *******
print('Running reconstruction using CDIModel.Adam_optimize on provided' +
' reconstruction_device,', reconstruction_device)
t.manual_seed(0)
# We only need to test the first loop to ensure it's identical
for i, iterations in enumerate(epoch_tup[:1]):
for loss in model.Adam_optimize(iterations,
dataset,
lr=lr_tup[i],
batch_size=batch_size_tup[i]):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# Ensure equivalency between the model reconstructions during the first
# pass, where they should be identical
assert np.allclose(model_recon.loss_history[:epoch_tup[0]], model.loss_history[:epoch_tup[0]])
# Ensure reconstructions have reached a certain loss tolerance. This just
# comes from running a reconstruction when it was working well and
# choosing a rough value. If it triggers this assertion error, something
# changed to make the final quality worse!
assert model_recon.loss_history[-1] < 0.0001
@pytest.mark.slow
def test_LBFGS_RPI(optical_data_ss_cxi,
optical_ptycho_incoherent_pickle,
reconstruction_device,
show_plot):
"""
This test checks out several things with the transmission RPI dataset
1) Calls to Reconstructor.adjust_optimizer is updating the
hyperparameters
2) Ensure `recon.model` points to the original `model`
3) Reconstructions performed by `LBFGS.optimize` and
`model.LBFGS_optimize` calls produce identical results when
run over one round of reconstruction.
4) The quality of the reconstruction remains below a specified
threshold.
5) Ensure that the RPI model works fine and dandy with the
Reconstructors.
"""
with open(optical_ptycho_incoherent_pickle, 'rb') as f:
ptycho_results = pickle.load(f)
probe = ptycho_results['probe']
background = ptycho_results['background']
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(optical_data_ss_cxi)
model = cdtools.models.RPI.from_dataset(dataset, probe, [500, 500],
background=background, n_modes=2,
initialization='random')
# Prepare two sets of models for the comparative reconstruction
model_recon = deepcopy(model)
model.to(device=reconstruction_device)
model_recon.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
# ******* Reconstructions with LBFGSReconstructor.optimize ******
print('Running reconstruction using LBFGSReconstructor.' +
'optimize on provided reconstruction_device,', reconstruction_device)
recon = cdtools.reconstructors.LBFGSReconstructor(model=model_recon,
dataset=dataset)
t.manual_seed(0)
# Run a reconstruction
reg_factor_tup = ([0.05, 0.05], [0.001, 0.1])
epoch_tup = (30, 50)
for i, iterations in enumerate(epoch_tup):
for loss in recon.optimize(iterations,
lr=0.4,
regularization_factor=reg_factor_tup[i]):
if show_plot and i == 0:
model_recon.inspect(dataset)
print(model_recon.report())
# Check hyperparameter update (or lack thereof)
assert recon.optimizer.param_groups[0]['lr'] == 0.4
if show_plot:
model_recon.inspect(dataset)
model_recon.compare(dataset)
# Check model pointing
assert id(model_recon) == id(recon.model)
# ******* Reconstructions with CDIModel.LBFGS_optimize ******
print('Running reconstruction using CDIModel.LBFGS_optimize.' +
'optimize on provided reconstruction_device,', reconstruction_device)
t.manual_seed(0)
for i, iterations in enumerate(epoch_tup[:1]):
for loss in model.LBFGS_optimize(iterations,
dataset,
lr=0.4,
regularization_factor=reg_factor_tup[i]): # noqa
if show_plot and i == 0:
model.inspect(dataset)
print(model.report())
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# Check loss equivalency between the two reconstructions
assert np.allclose(model.loss_history[:epoch_tup[0]], model_recon.loss_history[:epoch_tup[0]])
# The final loss when testing this was 2.28607e-3. Based on this, we set
# a threshold of 2.3e-3 for the tested loss. If this value has been
# exceeded, the reconstructions have gotten worse.
assert model_recon.loss_history[-1] < 0.0023
@pytest.mark.slow
def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
"""
This test checks out several things with the Au particle dataset
1) Calls to Reconstructor.adjust_optimizer is updating the
hyperparameters
3) Ensure `recon.model` points to the original `model`
4) Reconstructions performed by `SGD.optimize` and
`model.SGD_optimize` calls produce identical results
when run over one round of reconstruction.
5) The quality of the reconstruction remains below a specified
threshold.
5) Ensure that the FancyPtycho model works fine and dandy with the
Reconstructors.
The hyperparameters used in this test are not optimized to produce
a super-high-quality reconstruction. Instead, I just need A reconstruction
to do some kind of comparative assessment.
"""
print('\nTesting performance on the standard gold balls dataset ' +
'with reconstructors.SGDReconstructor')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
pad = 10
dataset.pad(pad)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
probe_support_radius=50,
propagation_distance=2e-6,
units='um',
probe_fourier_crop=pad
)
model.translation_offsets.data += 0.7 * \
t.randn_like(model.translation_offsets)
model.weights.requires_grad = False
# Make a copy of the model
model_recon = deepcopy(model)
model.to(device=reconstruction_device)
model_recon.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
# ******* Reconstructions with SGDReconstructor.optimize *******
print('Running reconstruction using SGDReconstructor.optimize' +
' on provided reconstruction_device,', reconstruction_device)
recon = cdtools.reconstructors.SGDReconstructor(model=model_recon,
dataset=dataset)
t.manual_seed(0)
# Run a reconstruction
epochs = 50
lr = 0.00000005
batch_size = 40
for loss in recon.optimize(epochs,
lr=lr,
batch_size=batch_size):
print(model_recon.report())
if show_plot and model_recon.epoch % 10 == 0:
model_recon.inspect(dataset)
# Check hyperparameter update
assert recon.optimizer.param_groups[0]['lr'] == lr
assert recon.data_loader.batch_size == batch_size
# Ensure that recon does not have sampler as an attribute (only used in
# multi-GPU)
assert not hasattr(recon, 'sampler')
# Ensure recon.model points to the original model
assert id(model_recon) == id(recon.model)
model_recon.tidy_probes()
if show_plot:
model_recon.inspect(dataset)
model_recon.compare(dataset)
# ******* Reconstructions with cdtools.CDIModel.SGD_optimize *******
print('Running reconstruction using CDIModel.SGD_optimize on provided' +
' reconstruction_device,', reconstruction_device)
t.manual_seed(0)
for loss in model.SGD_optimize(epochs,
dataset,
lr=lr,
batch_size=batch_size):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# Ensure equivalency between the model reconstructions
assert np.allclose(model_recon.loss_history[-1], model.loss_history[-1])
# The final loss when testing this was 7.12188e-4. Based on this, we set
# a threshold of 7.2e-4 for the tested loss. If this value has been
# exceeded, the reconstructions have gotten worse.
assert model.loss_history[-1] < 0.00072
tests/test_version.py
+18
View File
@@ -0,0 +1,18 @@
from cdtools import __version__
import re
def test_version_exists():
"""Test that version is defined and not empty."""
assert __version__
assert isinstance(__version__, str)
assert len(__version__) > 0
def test_version_format():
"""Test that version follows semantic versioning format."""
# Basic semantic versioning pattern (X.Y.Z with optional pre-release)
pattern = r"^\d+\.\d+\.\d+(?:[-.]?(?:alpha|beta|rc|dev)\d*)?$"
assert re.match(
pattern, __version__
), f"Version '{__version__}' doesn't follow semantic versioning"
tests/tools/test_analysis.py
+262 -219
View File
@@ -1,8 +1,7 @@
import numpy as np
import torch as t
from scipy import linalg as la
from scipy.sparse import linalg as spla
import torch as t
from itertools import combinations
from cdtools.tools import analysis, initializers
@@ -15,10 +14,10 @@ def test_product_svd():
A = np.random.rand(*shape_A) + 1j * np.random.rand(*shape_A)
B = np.random.rand(*shape_B) + 1j * np.random.rand(*shape_B)
AB = np.matmul(A,B)
AB = np.matmul(A, B)
U_1, S_1, Vh_1 = t.linalg.svd(t.as_tensor(AB), full_matrices=False)
U_2, S_2, Vh_2 = analysis.product_svd(t.as_tensor(A),t.as_tensor(B))
U_2, S_2, Vh_2 = analysis.product_svd(t.as_tensor(A), t.as_tensor(B))
check_AB = U_2 @ t.diag_embed(S_2).to(dtype=Vh_2.dtype) @ Vh_2
# So, at a minimum, U S Vh = AB
@@ -28,23 +27,23 @@ def test_product_svd():
# singular vector, so all we can ask for in the comparison is that the
# magnitudes here are
assert np.allclose(S_1[:rank].numpy(), S_2.numpy())
prod_U = U_1[:,:rank].transpose(0,1).conj() @ U_2
prod_Vh = Vh_1[:rank,:] @ Vh_2.transpose(0,1).conj()
prod_U = U_1[:, :rank].transpose(0, 1).conj() @ U_2
prod_Vh = Vh_1[:rank, :] @ Vh_2.transpose(0, 1).conj()
assert np.allclose(t.abs(prod_U).numpy(), np.eye(rank))
assert np.allclose(t.abs(prod_Vh).numpy(), np.eye(rank))
# Confirms that the phases are consistent between the two, I think
# it's redundant with the first check but I'm not sure
assert np.allclose(prod_Vh.numpy(), prod_U.numpy())
# test with numpy
U_3, S_3, Vh_3 = analysis.product_svd(A,B)
U_3, S_3, Vh_3 = analysis.product_svd(A, B)
assert isinstance(U_3, np.ndarray)
assert isinstance(S_3, np.ndarray)
assert isinstance(Vh_3, np.ndarray)
assert np.allclose(S_1[:rank].numpy(), S_3)
prod_U = U_1[:,:rank].transpose(0,1).numpy().conj() @ U_3
prod_Vh = Vh_1[:rank,:].numpy() @ Vh_3.transpose().conj()
prod_U = U_1[:, :rank].transpose(0, 1).numpy().conj() @ U_3
prod_Vh = Vh_1[:rank, :].numpy() @ Vh_3.transpose().conj()
assert np.allclose(np.abs(prod_U), np.eye(rank))
assert np.allclose(np.abs(prod_Vh), np.eye(rank))
# Confirms that the phases are consistent between the two, I think
@@ -55,42 +54,49 @@ def test_product_svd():
def test_orthogonalize_probes():
op = analysis.orthogonalize_probes
probe_xs = np.arange(64) - 32
probe_ys = np.arange(76) - 38
probe_Ys, probe_Xs = np.meshgrid(probe_ys, probe_xs)
probe_Rs = np.sqrt(probe_Xs**2 + probe_Ys**2)
probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2 + 1j)),
3*np.exp(-probe_Rs**2 / (2 * 12**2 - 3j)),
1*np.exp(-probe_Rs**2 / (2 * 15**2))])
probes = np.array(
[
10 * np.exp(-(probe_Rs**2) / (2 * 10**2 + 1j)),
3 * np.exp(-(probe_Rs**2) / (2 * 12**2 - 3j)),
1 * np.exp(-(probe_Rs**2) / (2 * 15**2)),
]
)
weight_matrix_none = None
weight_matrix_single = np.random.randn(1,3) + 1j * np.random.randn(1,3)
weight_matrix_small = np.random.randn(2,3) + 1j * np.random.randn(2,3)
weight_matrix_medium = np.random.randn(3,3) + 1j * np.random.randn(3,3)
weight_matrix_large = np.random.randn(7,3) + 1j * np.random.randn(7,3)
weight_matrix_single = np.random.randn(1, 3) + 1j * np.random.randn(1, 3)
weight_matrix_small = np.random.randn(2, 3) + 1j * np.random.randn(2, 3)
weight_matrix_medium = np.random.randn(3, 3) + 1j * np.random.randn(3, 3)
weight_matrix_large = np.random.randn(7, 3) + 1j * np.random.randn(7, 3)
weight_matrices = [
weight_matrix_none,
weight_matrix_single,
weight_matrix_small,
weight_matrix_medium,
weight_matrix_large
weight_matrix_large,
]
for weight_matrix in weight_matrices:
ortho_probes_np, rwm_np = op(probes, weight_matrix=weight_matrix,
return_reexpressed_weights=True)
ortho_probes_np, rwm_np = op(
probes, weight_matrix=weight_matrix, return_reexpressed_weights=True
)
assert isinstance(ortho_probes_np, np.ndarray)
assert isinstance(rwm_np, np.ndarray)
probes_t = t.as_tensor(probes)
wm_t = (t.as_tensor(weight_matrix) if weight_matrix is not None
else weight_matrix)
ortho_probes_t, rwm_t = op(probes_t, weight_matrix=wm_t,
return_reexpressed_weights=True)
probes_t = t.as_tensor(probes)
wm_t = (
t.as_tensor(weight_matrix) if weight_matrix is not None else weight_matrix
)
ortho_probes_t, rwm_t = op(
probes_t, weight_matrix=wm_t, return_reexpressed_weights=True
)
assert t.is_tensor(ortho_probes_t)
assert t.is_tensor(rwm_t)
@@ -105,15 +111,17 @@ def test_orthogonalize_probes():
realized_probes = probes
calculated_probes = np.tensordot(rwm_np, ortho_probes_np, axes=1)
assert np.allclose(realized_probes, calculated_probes)
# Now we test if the orthogonalized probes are orthogonalized
reshaped_probes = ortho_probes_np.reshape(
(ortho_probes_np.shape[0],
ortho_probes_np.shape[1] * ortho_probes_np.shape[2]))
products = np.matmul(reshaped_probes,
reshaped_probes.conj().transpose())
(
ortho_probes_np.shape[0],
ortho_probes_np.shape[1] * ortho_probes_np.shape[2],
)
)
products = np.matmul(reshaped_probes, reshaped_probes.conj().transpose())
if weight_matrix is not None:
output_nmodes = min(weight_matrix.shape[0], probes.shape[0])
@@ -121,23 +129,22 @@ def test_orthogonalize_probes():
output_nmodes = probes.shape[0]
for i in range(output_nmodes):
for j in range(i+1, output_nmodes):
assert np.isclose(products[i,j], 0)
for j in range(i + 1, output_nmodes):
assert np.isclose(products[i, j], 0)
# And now we test if they multiply to the same density matrix as
# the original probes + weight matrix
reshaped_realized_probes = realized_probes.reshape(
(realized_probes.shape[0],
realized_probes.shape[1] * realized_probes.shape[2]))
(
realized_probes.shape[0],
realized_probes.shape[1] * realized_probes.shape[2],
)
)
dm_original = np.matmul(
reshaped_realized_probes.conj().transpose(),
reshaped_realized_probes
)
dm_output = np.matmul(
reshaped_probes.conj().transpose(),
reshaped_probes
reshaped_realized_probes.conj().transpose(), reshaped_realized_probes
)
dm_output = np.matmul(reshaped_probes.conj().transpose(), reshaped_probes)
assert np.allclose(dm_original, dm_output)
# And finally, we confirm that what we have are the eigenvectors/values
@@ -150,53 +157,58 @@ def test_orthogonalize_probes():
# is undefined
assert np.allclose(np.abs(cross_products), np.abs(products))
def test_standardize():
# Start by making a probe and object that should meet the standardization
# conditions
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
probe = initializers.gaussian((230, 240), (20, 20), curvature=(0.01, 0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe) ** 2))
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.angle(np.sum(probe)) < 1e-7
assert np.isclose(1, np.sum(np.abs(probe) ** 2) / len(probe.ravel()))
assert np.angle(np.sum(probe)) < 2e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
(obj.shape[1]//8)*3:(obj.shape[1]//8)*5]
obj = 30 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5))
obj_slice = np.s_[
(obj.shape[0] // 8) * 3:(obj.shape[0] // 8) * 5,
(obj.shape[1] // 8) * 3:(obj.shape[1] // 8) * 5,
]
obj = obj * np.exp(-1j * np.angle(np.sum(obj[obj_slice])))
assert np.isclose(0,np.angle(np.sum(obj[obj_slice])))
assert np.isclose(0, np.angle(np.sum(obj[obj_slice])))
# Then make a nonstandard version of them and standardize it
# First, don't add a phase ramp and test
test_probe = probe * 37.6 * np.exp(1j*0.35)
test_obj = obj / 37.6 * np.exp(1j*1.43)
test_probe = probe * 37.6 * np.exp(1j * 0.35)
test_obj = obj / 37.6 * np.exp(1j * 1.43)
s_probe, s_obj = analysis.standardize(test_probe, test_obj)
assert np.allclose(probe, s_probe)
assert np.allclose(obj, s_obj)
# Test that it works on torch tensors
s_probe, s_obj = analysis.standardize(t.as_tensor(test_probe,dtype=t.complex64), t.as_tensor(test_obj,dtype=t.complex64))
s_probe, s_obj = analysis.standardize(
t.as_tensor(test_probe, dtype=t.complex64),
t.as_tensor(test_obj, dtype=t.complex64),
)
s_probe = s_probe.numpy()
s_obj = s_obj.numpy()
assert np.allclose(probe, s_probe)
assert np.allclose(obj, s_obj)
# Then do one with a phase ramp
phase_ramp_dir = (np.random.rand(2) - 0.5)
phase_ramp_dir = np.random.rand(2) - 0.5
probe_Xs, probe_Ys = np.mgrid[:probe.shape[0],:probe.shape[1]]
phase_ramp = np.exp(1j*probe_Ys * phase_ramp_dir[1]+
1j*probe_Xs * phase_ramp_dir[0])
probe_Xs, probe_Ys = np.mgrid[: probe.shape[0], : probe.shape[1]]
phase_ramp = np.exp(
1j * probe_Ys * phase_ramp_dir[1] + 1j * probe_Xs * phase_ramp_dir[0]
)
test_probe = test_probe * phase_ramp
obj_Xs, obj_Ys = np.mgrid[:obj.shape[0],:obj.shape[1]]
obj_phase_ramp = np.exp(-1j*obj_Ys * phase_ramp_dir[1]+
-1j*obj_Xs * phase_ramp_dir[0])
obj_Xs, obj_Ys = np.mgrid[: obj.shape[0], : obj.shape[1]]
obj_phase_ramp = np.exp(
-1j * obj_Ys * phase_ramp_dir[1] + -1j * obj_Xs * phase_ramp_dir[0]
)
test_obj = test_obj * obj_phase_ramp
s_probe, s_obj = analysis.standardize(test_probe, test_obj, correct_ramp=True)
@@ -205,12 +217,12 @@ def test_standardize():
assert np.max(s_obj - obj) / np.max(np.abs(obj)) < 1e-4
# Finally a test with the phase ramp and multiple probes
subdominant_probe = 0.1*np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
subdominant_probe = (0.1 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5)))
subdominant_probe = subdominant_probe * np.exp(-1j * np.angle(np.sum(subdominant_probe)))
test_subdominant_probe = subdominant_probe * 37.6
test_subdominant_probe = test_subdominant_probe * phase_ramp
incoh_probe = np.array([test_probe,test_subdominant_probe])
incoh_probe = np.array([test_probe, test_subdominant_probe])
s_probe, s_obj = analysis.standardize(incoh_probe, test_obj, correct_ramp=True)
@@ -219,7 +231,6 @@ def test_standardize():
assert np.max(s_probe[1] - subdominant_probe) / np.max(np.abs(subdominant_probe)) < 1e-4
def test_synthesize_reconstructions():
# I can only really test for a lack of failures, so I think my plan
# will be to create a dataset that just needs to be added and see that
@@ -227,140 +238,153 @@ def test_synthesize_reconstructions():
# Start by making a probe and object that should meet the standardization
# conditions
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
probe = initializers.gaussian((230, 240), (20, 20), curvature=(0.01, 0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe) ** 2))
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.abs(np.angle(np.sum(probe))) < 1e-7
assert np.isclose(1, np.sum(np.abs(probe) ** 2) / len(probe.ravel()))
assert np.abs(np.angle(np.sum(probe))) < 2e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
(obj.shape[1]//8)*3:(obj.shape[1]//8)*5]
obj = 30 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5))
obj_slice = np.s_[
(obj.shape[0] // 8) * 3:(obj.shape[0] // 8) * 5,
(obj.shape[1] // 8) * 3:(obj.shape[1] // 8) * 5,
]
obj = obj * np.exp(-1j * np.angle(np.sum(obj[obj_slice])))
assert np.isclose(0,np.angle(np.sum(obj[obj_slice])))
assert np.isclose(0, np.angle(np.sum(obj[obj_slice])))
# Now I make stacks of identical probes and objects
probes = [probe,probe,probe,probe]
probes = [probe, probe, probe, probe]
probe = np.copy(probe)
objects = [obj,obj,obj,obj]
objects = [obj, obj, obj, obj]
obj = np.copy(obj)
s_probe, s_obj, obj_stack = analysis.synthesize_reconstructions(probes,objects)
s_probe, s_obj, obj_stack = analysis.synthesize_reconstructions(probes, objects)
assert np.max(s_probe - probe) < 2e-5
assert np.max(s_obj - obj) < 2e-5
for t_obj in obj_stack:
assert np.max(t_obj - obj) < 5e-5
def test_calc_consistency_prtf():
# Create an object with a specific structure
obj = 30 * np.random.rand(1030,1040) * np.exp(1j * (np.random.rand(1030,1040) - 0.5))
obj = (30 * np.random.rand(1030, 1040) * np.exp(1j * (np.random.rand(1030, 1040) - 0.5)))
#
synth_obj = np.sqrt(0.7) * obj
obj_stack = [obj,obj,obj,obj]
obj_stack = [obj, obj, obj, obj]
basis = np.array([[0, 2, 0], [3, 0, 0]])
basis = np.array([[0,2,0],
[3,0,0]])
freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis)
assert np.allclose(prtf, 0.7)
freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis, nbins=30)
assert np.allclose(prtf, 0.7)
# Check that it also works with torch input
t_synth_obj = t.as_tensor(synth_obj)
t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, basis, nbins=30)
freqs, prtf = analysis.calc_consistency_prtf(
t_synth_obj, t_obj_stack, basis, nbins=30
)
assert np.allclose(prtf.numpy(), 0.7)
# And also when the basis is in torch
t_synth_obj = t.as_tensor(synth_obj)
t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30)
freqs, prtf = analysis.calc_consistency_prtf(
t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30
)
assert np.allclose(prtf.numpy(), 0.7)
# Check that is uses the right number of bins
assert len(prtf) == 30
assert len(freqs) == 30
# Check that the maximum frequency is correct for the basis
assert np.isclose(freqs[-1]-freqs[-2] + freqs[-1], np.sqrt(1/4**2 + 1/6**2))
assert np.isclose(freqs[-1] - freqs[-2] + freqs[-1], np.sqrt(1 / 4**2 + 1 / 6**2))
def test_calc_deconvolved_cross_correlation():
obj1 = np.random.rand(200,300) + 1j * np.random.rand(200,300)
obj2 = np.random.rand(200,300) + 1j * np.random.rand(200,300)
obj1 = np.random.rand(200, 300) + 1j * np.random.rand(200, 300)
obj2 = np.random.rand(200, 300) + 1j * np.random.rand(200, 300)
cor_fft = np.fft.fft2(obj1) * np.conj(np.fft.fft2(obj2))
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
np_cor = np.fft.ifft2(cor_fft / np.abs(cor_fft))
# test with numpy inputs
test_cor = analysis.calc_deconvolved_cross_correlation(obj1,obj2, im_slice=np.s_[:,:])
test_cor = analysis.calc_deconvolved_cross_correlation(
obj1, obj2, im_slice=np.s_[:, :]
)
assert np.allclose(test_cor, np_cor)
# test with pytorch inputs
obj1_t = t.as_tensor(obj1)
obj2_t = t.as_tensor(obj2)
test_cor_t = analysis.calc_deconvolved_cross_correlation(obj1_t,obj2_t, im_slice=np.s_[:,:])
test_cor_t = analysis.calc_deconvolved_cross_correlation(
obj1_t, obj2_t, im_slice=np.s_[:, :]
)
assert np.allclose(test_cor_t.numpy(), np_cor)
def test_calc_frc():
obj1 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
obj2 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
obj1 = np.random.rand(270, 230) + 1j * np.random.rand(270, 230)
obj2 = np.random.rand(270, 230) + 1j * np.random.rand(270, 230)
basis = np.array([[0,2,0],
[3,0,0]])
basis = np.array([[0, 2, 0], [3, 0, 0]])
nbins = 100
snr = 2
cor_fft = np.fft.fftshift(np.fft.fft2(obj1[10:-10,20:-20])) * \
np.fft.fftshift(np.conj(np.fft.fft2(obj2[10:-10,20:-20])))
F1 = np.abs(np.fft.fftshift(np.fft.fft2(obj1[10:-10,20:-20])))**2
F2 = np.abs(np.fft.fftshift(np.fft.fft2(obj2[10:-10,20:-20])))**2
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[0],d=di))
j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
cor_fft = np.fft.fftshift(np.fft.fft2(obj1[10:-10, 20:-20])) * np.fft.fftshift(
np.conj(np.fft.fft2(obj2[10:-10, 20:-20]))
)
numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft)
denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1)
denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2)
n_pix, bins = np.histogram(Rs,bins=nbins)
F1 = np.abs(np.fft.fftshift(np.fft.fft2(obj1[10:-10, 20:-20]))) ** 2
F2 = np.abs(np.fft.fftshift(np.fft.fft2(obj2[10:-10, 20:-20]))) ** 2
di = np.linalg.norm(basis[:, 0])
dj = np.linalg.norm(basis[:, 1])
i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[0], d=di))
j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[1], d=dj))
Js, Is = np.meshgrid(j_freqs, i_freqs)
Rs = np.sqrt(Is**2 + Js**2)
numerator, bins = np.histogram(Rs, bins=nbins, weights=cor_fft)
denominator_F1, bins = np.histogram(Rs, bins=nbins, weights=F1)
denominator_F2, bins = np.histogram(Rs, bins=nbins, weights=F2)
n_pix, bins = np.histogram(Rs, bins=nbins)
bins = bins[:-1]
frc = numerator / np.sqrt(denominator_F1*denominator_F2)
# This moves from combined-image SNR to single-image SNR
frc = numerator / np.sqrt(denominator_F1 * denominator_F2)
# This moves from combined-image SNR to single-image SNR
snr /= 2
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
(1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / (
1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix)
)
test_bins, test_frc, test_threshold = analysis.calc_frc(
obj1, obj2, basis, im_slice=np.s_[10:-10,20:-20],
nbins=100, snr=2, limit='corner')
obj1,
obj2,
basis,
im_slice=np.s_[10:-10, 20:-20],
nbins=100,
snr=2,
limit="corner",
)
assert np.allclose(bins, test_bins)
assert np.allclose(frc, test_frc)
assert np.allclose(threshold, test_threshold)
@@ -374,33 +398,40 @@ def test_calc_frc():
obj1_torch,
obj2_torch,
basis_torch,
im_slice=np.s_[10:-10,20:-20], nbins=100, snr=2, limit='corner')
im_slice=np.s_[10:-10, 20:-20],
nbins=100,
snr=2,
limit="corner",
)
assert np.allclose(bins, test_bins_t.numpy())
assert np.allclose(frc, test_frc_t.numpy())
assert np.allclose(threshold, test_threshold_t.numpy())
def test_calc_rms_error():
field_1 = t.rand(14,19, dtype=t.complex64)
field_2 = t.rand(14,19, dtype=t.complex64)
field_1 = t.rand(14, 19, dtype=t.complex64)
field_2 = t.rand(14, 19, dtype=t.complex64)
# Check that the calculation is insensitive to phase
assert t.allclose(analysis.calc_rms_error(field_1, field_2),
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2))
assert t.allclose(
analysis.calc_rms_error(field_1, field_2),
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2),
)
# And that it is sensitive to phase if we turn off the
# And that it is sensitive to phase if we turn off the
assert not t.allclose(
analysis.calc_rms_error(field_1, field_2, align_phases=False),
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2,
align_phases=False))
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2, align_phases=False),
)
# Check that the result is positive
assert analysis.calc_rms_error(field_1, field_2) > 0
# And that it is a smaller number with align_phases on
assert (analysis.calc_rms_error(field_1, field_2) <=
analysis.calc_rms_error(field_1, field_2, align_phases=False))
assert analysis.calc_rms_error(field_1, field_2) <= analysis.calc_rms_error(
field_1, field_2, align_phases=False
)
# Now we check against an explicit implementation:
gamma = field_1 * t.conj(field_2)
@@ -408,114 +439,126 @@ def test_calc_rms_error():
# This is an alternate way of doing the calculation. Actually, would this
# be a better implementation anyway? Probably no difference tbh.
rms_error_nophase = t.sqrt((t.mean(t.abs(field_1)**2) +
t.mean(t.abs(field_2)**2) -
2 * t.abs(t.mean(field_1 * t.conj(field_2)))))
assert t.allclose(rms_error_nophase,
analysis.calc_rms_error(field_1, field_2))
rms_error_nophase = t.sqrt(
(
t.mean(t.abs(field_1) ** 2)
+ t.mean(t.abs(field_2) ** 2)
- 2 * t.abs(t.mean(field_1 * t.conj(field_2)))
)
)
assert t.allclose(rms_error_nophase, analysis.calc_rms_error(field_1, field_2))
rms_error_phase = t.sqrt((t.mean(t.abs(field_1)**2) +
t.mean(t.abs(field_2)**2) -
2 * t.real(t.mean(field_1 * t.conj(field_2)))))
rms_error_phase = t.sqrt(
(
t.mean(t.abs(field_1) ** 2)
+ t.mean(t.abs(field_2) ** 2)
- 2 * t.real(t.mean(field_1 * t.conj(field_2)))
)
)
assert t.allclose(rms_error_phase,
analysis.calc_rms_error(field_1, field_2,
align_phases=False))
assert t.allclose(
rms_error_phase, analysis.calc_rms_error(field_1, field_2, align_phases=False)
)
# Now let's test that it works along a dimension:
field_1 = t.rand(3,14,19, dtype=t.complex64)
field_2 = t.rand(3,14,19, dtype=t.complex64)
field_1 = t.rand(3, 14, 19, dtype=t.complex64)
field_2 = t.rand(3, 14, 19, dtype=t.complex64)
result = analysis.calc_rms_error(field_1, field_2, normalize=True)
assert (result.shape == t.Size([3]))
assert result.shape == t.Size([3])
for i in range(3):
assert t.allclose(analysis.calc_rms_error(field_1[i],
field_2[i],
normalize=True),
result[i])
assert t.allclose(
analysis.calc_rms_error(field_1[i], field_2[i], normalize=True), result[i]
)
def test_calc_fidelity():
fields_1 = t.rand(2,30,17, dtype=t.complex128)
fields_2 = t.rand(3,30,17, dtype=t.complex128)
fields_1 = t.rand(2, 30, 17, dtype=t.complex128)
fields_2 = t.rand(3, 30, 17, dtype=t.complex128)
dm_1 = t.reshape(fields_1, (2,-1))
dm_1 = t.tensordot(dm_1.transpose(0,1), dm_1.conj(), dims=1).numpy()
dm_2 = t.reshape(fields_2, (3,-1))
dm_2 = t.tensordot(dm_2.transpose(0,1), dm_2.conj(), dims=1).numpy()
dm_1 = t.reshape(fields_1, (2, -1))
dm_1 = t.tensordot(dm_1.transpose(0, 1), dm_1.conj(), dims=1).numpy()
dm_2 = t.reshape(fields_2, (3, -1))
dm_2 = t.tensordot(dm_2.transpose(0, 1), dm_2.conj(), dims=1).numpy()
sqrt_dm_1 = la.sqrtm(dm_1).astype(dm_1.dtype)
inner_mat = la.sqrtm(np.dot(np.dot(sqrt_dm_1,dm_2), sqrt_dm_1))
inner_mat = inner_mat.astype(dm_1.dtype) #la.sqrtm doubles the precision
fidelity = t.as_tensor(np.abs(np.trace(inner_mat))**2)
inner_mat = la.sqrtm(np.dot(np.dot(sqrt_dm_1, dm_2), sqrt_dm_1))
inner_mat = inner_mat.astype(dm_1.dtype) # la.sqrtm doubles the precision
fidelity = t.as_tensor(np.abs(np.trace(inner_mat)) ** 2)
assert t.isclose(fidelity, analysis.calc_fidelity(fields_1, fields_2))
# Check that it reduces to the overlap for coherent fields
fields_1 = t.rand(1,30,17, dtype=t.complex128)
fields_2 = t.rand(1,30,17, dtype=t.complex128)
fields_1 = t.rand(1, 30, 17, dtype=t.complex128)
fields_2 = t.rand(1, 30, 17, dtype=t.complex128)
assert t.isclose(t.abs(t.sum(fields_1*fields_2.conj()))**2,
analysis.calc_fidelity(fields_1, fields_2))
assert t.isclose(
t.abs(t.sum(fields_1 * fields_2.conj())) ** 2,
analysis.calc_fidelity(fields_1, fields_2),
)
# Checking that it works with extra dimensions
fields_1 = t.rand(3,3,30,17, dtype=t.complex128)
fields_2 = t.rand(3,1,30,17, dtype=t.complex128)
field_3 = t.rand(1,30,17, dtype=t.complex128)
fields_1 = t.rand(3, 3, 30, 17, dtype=t.complex128)
fields_2 = t.rand(3, 1, 30, 17, dtype=t.complex128)
field_3 = t.rand(1, 30, 17, dtype=t.complex128)
fidelities = analysis.calc_fidelity(fields_1, fields_2)
fidelities_2 = analysis.calc_fidelity(fields_1, field_3)
for i in range(3):
assert t.isclose(analysis.calc_fidelity(fields_1[i], fields_2[i]),
fidelities[i])
assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3),
fidelities_2[i])
assert t.isclose(
analysis.calc_fidelity(fields_1[i], fields_2[i]), fidelities[i]
)
assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3), fidelities_2[i])
# Check that the diensionality argument works
fields_1 = t.rand(3,2,12, dtype=t.complex128)
fields_2 = t.rand(3,2,12, dtype=t.complex128)
fields_1 = t.rand(3, 2, 12, dtype=t.complex128)
fields_2 = t.rand(3, 2, 12, dtype=t.complex128)
assert (analysis.calc_fidelity(fields_1, fields_2, dims=1).shape
== t.Size([3]))
fields_1 = t.rand(3,2,12,4,5, dtype=t.complex128)
fields_2 = t.rand(3,2,12,4,5, dtype=t.complex128)
assert analysis.calc_fidelity(fields_1, fields_2, dims=1).shape == t.Size([3])
fields_1 = t.rand(3, 2, 12, 4, 5, dtype=t.complex128)
fields_2 = t.rand(3, 2, 12, 4, 5, dtype=t.complex128)
assert analysis.calc_fidelity(fields_1, fields_2, dims=3).shape == t.Size([3])
assert (analysis.calc_fidelity(fields_1, fields_2, dims=3).shape
== t.Size([3]))
def test_calc_generalized_rms_error():
# Test that it matches the rms error for coherent fields
fields_1 = t.rand(1,30,17, dtype=t.complex128)
fields_2 = t.rand(1,30,17, dtype=t.complex128)
assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_rms_error(fields_1[0], fields_2[0],
align_phases=True))
fields_1 = t.rand(1, 30, 17, dtype=t.complex128)
fields_2 = t.rand(1, 30, 17, dtype=t.complex128)
assert t.isclose(
analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_rms_error(fields_1[0], fields_2[0], align_phases=True),
)
# Test that it is independent of field order
fields_1 = t.rand(5,30,17, dtype=t.complex128)
fields_2 = t.rand(3,30,17, dtype=t.complex128)
fields_1 = t.rand(5, 30, 17, dtype=t.complex128)
fields_2 = t.rand(3, 30, 17, dtype=t.complex128)
fields_3 = fields_2.flip(0)
assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_generalized_rms_error(fields_1, fields_3))
assert t.isclose(
analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_generalized_rms_error(fields_1, fields_3),
)
# Test with leading dimensions
fields_1 = t.rand(3,4,2,10,17, dtype=t.complex128)
fields_2 = t.rand(3,4,3,10,17, dtype=t.complex128)
assert (analysis.calc_generalized_rms_error(fields_1, fields_2).shape
== t.Size([3,4]))
fields_1 = t.rand(3, 4, 2, 10, 17, dtype=t.complex128)
fields_2 = t.rand(3, 4, 3, 10, 17, dtype=t.complex128)
assert analysis.calc_generalized_rms_error(fields_1, fields_2).shape == t.Size(
[3, 4]
)
# And test with different number of dimensions dims
# Test that it is independent of field order
fields_1 = t.rand(3,6,17, dtype=t.complex128)
fields_2 = t.rand(3,1,17, dtype=t.complex128)
fields_1 = t.rand(3, 6, 17, dtype=t.complex128)
fields_2 = t.rand(3, 1, 17, dtype=t.complex128)
fields_3 = fields_2.flip(0)
assert (analysis.calc_generalized_rms_error(fields_1, fields_2, dims=1).shape == t.Size([3]))
assert analysis.calc_generalized_rms_error(
fields_1, fields_2, dims=1
).shape == t.Size([3])
tests/tools/test_data.py
+93 -84
View File
@@ -1,66 +1,69 @@
from cdtools.tools import data
import numpy as np
import torch as t
import h5py
import pytest
import os
import datetime
import numbers
import pathlib
import h5py
import numpy as np
import torch as t
from cdtools.tools import data
#
# We start with a bunch of tests of the data loading capabilities
#
def test_get_entry_info(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
entry_info = data.get_entry_info(cxi)
for key in expected['entry metadata']:
assert entry_info[key] == expected['entry metadata'][key]
def test_get_sample_info(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
sample_info = data.get_sample_info(cxi)
if sample_info is None and \
('sample info' not in expected or
expected['sample info'] is None):
('sample info' not in expected or expected['sample info'] is None):
# Valid if no sample info is defined at all
continue
for key in expected['sample info']:
if isinstance(expected['sample info'][key],np.ndarray):
if isinstance(expected['sample info'][key], np.ndarray):
assert np.allclose(sample_info[key],
expected['sample info'][key])
else:
assert sample_info[key] == expected['sample info'][key]
def test_get_wavelength(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
assert np.isclose(expected['wavelength'],data.get_wavelength(cxi))
assert np.isclose(expected['wavelength'], data.get_wavelength(cxi))
def test_get_detector_geometry(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
distance, basis, corner = data.get_detector_geometry(cxi)
assert np.isclose(distance,expected['detector']['distance'])
assert np.allclose(basis,expected['detector']['basis'])
assert np.isclose(distance, expected['detector']['distance'])
assert np.allclose(basis, expected['detector']['basis'])
if isinstance(expected['detector']['corner'], np.ndarray):
assert np.allclose(corner, expected['detector']['corner'])
else:
assert corner == expected['detector']['corner']
def test_get_mask(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
mask = data.get_mask(cxi)
if expected['mask'] is None and mask is None:
continue
assert np.all(data.get_mask(cxi) == expected['mask'])
assert np.all(mask == expected['mask'])
def test_get_qe_mask(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
qe_mask = data.get_qe_mask(cxi)
if expected['qe_mask'] is None and qe_mask is None:
continue
assert np.allclose(qe_mask, expected['qe_mask'])
def test_get_dark(test_ptycho_cxis):
@@ -70,8 +73,8 @@ def test_get_dark(test_ptycho_cxis):
assert expected['dark'] is None
else:
assert np.allclose(dark, expected['dark'])
def test_get_data(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
patterns, axes = data.get_data(cxi)
@@ -91,8 +94,6 @@ def test_get_ptycho_translations(test_ptycho_cxis):
assert np.allclose(data.get_ptycho_translations(cxi),
expected['translations'])
#
# Then, write a test for the data saving. It should create a .cxi file
# using the data seving tools, and then check that when read with the
@@ -102,13 +103,13 @@ def test_get_ptycho_translations(test_ptycho_cxis):
def test_create_cxi(tmp_path):
data.create_cxi(tmp_path / 'test_create.cxi')
with h5py.File(tmp_path / 'test_create.cxi','r') as f:
with h5py.File(tmp_path / 'test_create.cxi', 'r') as f:
assert f['cxi_version'][()] == 160
assert 'entry_1' in f
def test_add_entry_info(tmp_path):
entry_info = {'experiment_identifier':'test of cxi file writing tools',
entry_info = {'experiment_identifier': 'test of cxi file writing tools',
'title': 'my cool experiment',
'start_time': datetime.datetime.now(),
'end_time': datetime.datetime.now()}
@@ -116,12 +117,11 @@ def test_add_entry_info(tmp_path):
with data.create_cxi(tmp_path / 'test_add_entry_info.cxi') as f:
data.add_entry_info(f, entry_info)
with h5py.File(tmp_path / 'test_add_entry_info.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_entry_info.cxi', 'r') as f:
read_entry_info = data.get_entry_info(f)
print(read_entry_info)
for key in entry_info:
if isinstance(entry_info[key], np.ndarray):
assert np.allclose(entry_info[key], read_entry_info[key])
@@ -130,18 +130,18 @@ def test_add_entry_info(tmp_path):
def test_add_sample_info(tmp_path):
sample_info = {'name':'A nice fake sample',
sample_info = {'name': 'A nice fake sample',
'concentration': 10,
'mass': 5.3,
'temperature': 76,
'description': 'A very nice sample',
'unit_cell': np.array([1,1,1,90.,90.,90.])}
'unit_cell': np.array([1, 1, 1, 90., 90., 90.])}
with data.create_cxi(tmp_path / 'test_add_sample_info.cxi') as f:
data.add_sample_info(f, sample_info)
with h5py.File(tmp_path / 'test_add_sample_info.cxi','r') as f:
read_sample_info = data.get_sample_info(f)
with h5py.File(tmp_path / 'test_add_sample_info.cxi', 'r') as f:
read_sample_info = data.get_sample_info(f)
for key in sample_info:
if isinstance(sample_info[key], np.ndarray):
@@ -150,7 +150,7 @@ def test_add_sample_info(tmp_path):
assert np.isclose(sample_info[key], read_sample_info[key])
else:
assert sample_info[key] == read_sample_info[key]
def test_add_source(tmp_path):
wavelength = 1e-9
@@ -159,26 +159,26 @@ def test_add_source(tmp_path):
with data.create_cxi(tmp_path / 'test_add_source.cxi') as f:
data.add_source(f, wavelength)
with h5py.File(tmp_path / 'test_add_source.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_source.cxi', 'r') as f:
# Check this directly since we want to make sure it saved
# the wavelength and energy
read_wavelength = f['entry_1/instrument_1/source_1/wavelength'][()]
read_energy = f['entry_1/instrument_1/source_1/energy'][()]
assert np.isclose( wavelength, read_wavelength)
assert np.isclose( energy, read_energy)
assert np.isclose(wavelength, read_wavelength)
assert np.isclose(energy, read_energy)
def test_add_detector(tmp_path):
distance = 0.34
basis = np.array([[0,-30e-6,0],
[-20e-6,0,0]]).astype(np.float32).transpose()
corner = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
basis = np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).astype(np.float32).transpose()
corner = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
with data.create_cxi(tmp_path / 'test_add_detector.cxi') as f:
data.add_detector(f, distance, basis, corner=corner)
with h5py.File(tmp_path / 'test_add_detector.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_detector.cxi', 'r') as f:
# Check this directly since we want to make sure it saved
# the pixel sizes
d1 = f['entry_1/instrument_1/detector_1']
@@ -190,100 +190,110 @@ def test_add_detector(tmp_path):
assert np.isclose(distance, read_distance)
assert np.allclose(basis, read_basis)
assert np.isclose(np.linalg.norm(basis[:,1]), read_x_pix)
assert np.isclose(np.linalg.norm(basis[:,0]), read_y_pix)
assert np.allclose(corner,read_corner)
assert np.isclose(np.linalg.norm(basis[:, 1]), read_x_pix)
assert np.isclose(np.linalg.norm(basis[:, 0]), read_y_pix)
assert np.allclose(corner, read_corner)
def test_add_mask(tmp_path):
mask = (np.random.rand(350,600) > 0.1).astype(np.uint8)
mask = (np.random.rand(350, 600) > 0.1).astype(np.uint8)
with data.create_cxi(tmp_path / 'test_add_mask.cxi') as f:
data.add_mask(f, mask)
with h5py.File(tmp_path / 'test_add_mask.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_mask.cxi', 'r') as f:
read_mask = data.get_mask(f)
assert np.all(mask == read_mask)
def test_add_qe_mask(tmp_path):
qe_mask = np.random.rand(350, 199).astype(np.float32)
with data.create_cxi(tmp_path / 'test_add_qe_mask.cxi') as f:
data.add_qe_mask(f, qe_mask)
with h5py.File(tmp_path / 'test_add_qe_mask.cxi', 'r') as f:
read_qe_mask = data.get_qe_mask(f)
assert np.allclose(qe_mask, read_qe_mask)
def test_add_dark(tmp_path):
dark = np.random.rand(350,620)
dark = np.random.rand(350, 620)
with data.create_cxi(tmp_path / 'test_add_dark.cxi') as f:
data.add_dark(f, dark)
with h5py.File(tmp_path / 'test_add_dark.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_dark.cxi', 'r') as f:
read_dark = data.get_dark(f)
print(dark.shape)
assert np.allclose(dark, read_dark)
def test_add_data(tmp_path):
# First test from numpy, with axes
fake_data = np.random.rand(100,256,256)
axes = ['translation','y','x']
fake_data = np.random.rand(100, 256, 256)
axes = ['translation', 'y', 'x']
with data.create_cxi(tmp_path / 'test_add_data.cxi') as f:
data.add_data(f, fake_data, axes)
with h5py.File(tmp_path / 'test_add_data.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_data.cxi', 'r') as f:
# Check this directly since we want to make sure it saved
# it in all the places it should have
read_data_1 = f['entry_1/data_1/data'][()]
read_data_2 = f['entry_1/instrument_1/detector_1/data'][()]
read_axes = str(f['entry_1/instrument_1/detector_1/data'].attrs['axes'].decode())
assert np.allclose(fake_data, read_data_1)
assert np.allclose(fake_data, read_data_2)
assert 'translation:y:x' == read_axes
# Then test from torch, without axes
fake_data = t.from_numpy(fake_data)
with data.create_cxi(tmp_path / 'test_add_data_torch.cxi') as f:
data.add_data(f, fake_data)
with h5py.File(tmp_path / 'test_add_data_torch.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_data_torch.cxi', 'r') as f:
read_data, axes = data.get_data(f)
assert np.allclose(fake_data.numpy(),read_data)
assert np.allclose(fake_data.numpy(), read_data)
def test_add_shot_to_shot_info(tmp_path):
analyzer = np.random.rand(100)
with data.create_cxi(tmp_path / 'test_add_shot_to_shot_info.cxi') as f:
data.add_shot_to_shot_info(f, analyzer, 'analyzer_angle')
with h5py.File(tmp_path / 'test_add_shot_to_shot_info.cxi') as f:
# Check this directly since we want to make sure it saved
# it in all the places it should have
read_analyzer_1 = f['entry_1/data_1/analyzer_angle'][()]
read_analyzer_2 = \
f['entry_1/instrument_1/detector_1/analyzer_angle'][()]
read_analyzer_2 = f['entry_1/instrument_1/detector_1/analyzer_angle'][()]
read_analyzer_3 = f['entry_1/sample_1/geometry_1/analyzer_angle'][()]
assert np.allclose(analyzer, read_analyzer_1)
assert np.allclose(analyzer, read_analyzer_2)
assert np.allclose(analyzer, read_analyzer_3)
def test_add_ptycho_translations(tmp_path):
translations = np.random.rand(3,100)
translations = np.random.rand(3, 100)
with data.create_cxi(tmp_path / 'test_add_ptycho_translations.cxi') as f:
data.add_ptycho_translations(f, translations)
with h5py.File(tmp_path / 'test_add_ptycho_translations.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_ptycho_translations.cxi', 'r') as f:
# Check this directly since we want to make sure it saved
# it in all the places it should have
read_translations_1 = f['entry_1/data_1/translation'][()]
read_translations_2 = \
f['entry_1/instrument_1/detector_1/translation'][()]
read_translations_2 = f['entry_1/instrument_1/detector_1/translation'][()]
read_translations_3 = f['entry_1/sample_1/geometry_1/translation'][()]
assert np.allclose(-translations, read_translations_1)
@@ -292,8 +302,7 @@ def test_add_ptycho_translations(tmp_path):
def test_nested_dict_to_h5(tmp_path, example_nested_dicts):
### Tests both nested_dict_to_h5 and h5_to_nested_dict
# Tests both nested_dict_to_h5 and h5_to_nested_dict
def check_dict_equality(truth, to_test):
for key in truth.keys():
if isinstance(truth[key], dict):
@@ -308,23 +317,24 @@ def test_nested_dict_to_h5(tmp_path, example_nested_dicts):
assert truth[key] == to_test[key]
else:
assert 0
for test_dict in example_nested_dicts:
filename = tmp_path / 'example_dataset.h5'
data.nested_dict_to_h5(filename, test_dict)
roundtrip = data.h5_to_nested_dict(filename)
check_dict_equality(test_dict, roundtrip)
def test_h5_to_nested_dict(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
# Just test that it runs without errors for these ones.
# A round-trip test is in test_nested_dict_to_h5
d = data.h5_to_nested_dict(cxi)
data.h5_to_nested_dict(cxi)
def test_nested_dict_to_numpy(example_nested_dicts):
def check_dict_numpyness(truth, to_test):
def check_dict_numpyness(truth, to_test):
for key in truth.keys():
if isinstance(truth[key], dict):
check_dict_numpyness(truth[key], to_test[key])
@@ -341,12 +351,11 @@ def test_nested_dict_to_numpy(example_nested_dicts):
for test_dict in example_nested_dicts:
numpy_dict = data.nested_dict_to_numpy(test_dict)
check_dict_numpyness(test_dict, numpy_dict)
def test_nested_dict_to_torch(example_nested_dicts):
check_dict_numpyness(test_dict, numpy_dict)
def check_dict_torchiness(truth, to_test):
def test_nested_dict_to_torch(example_nested_dicts):
def check_dict_torchiness(truth, to_test):
for key in truth.keys():
if isinstance(truth[key], dict):
check_dict_torchiness(truth[key], to_test[key])
tests/tools/test_image_processing.py
+57 -58
View File
@@ -1,18 +1,19 @@
import numpy as np
import torch as t
from scipy import ndimage
from cdtools.tools import image_processing, interactions
from scipy import ndimage
def test_centroid():
# Test single im
im = t.rand((30,40))
im = t.rand((30, 40))
sp_centroid = ndimage.center_of_mass(im.numpy())
centroid = image_processing.centroid(im)
assert t.allclose(centroid, t.Tensor(sp_centroid))
# Test stack o' ims
ims = t.rand((5,30,40))
ims = t.rand((5, 30, 40))
sp_centroids = [ndimage.center_of_mass(im.numpy())
for im in ims]
centroids = image_processing.centroid(ims)
@@ -21,33 +22,33 @@ def test_centroid():
def test_centroid_sq():
# Test single im
im = t.rand((30,40))
im = t.rand((30, 40))
sp_centroid = ndimage.center_of_mass(im.numpy()**2)
centroid = image_processing.centroid_sq(im)
assert t.allclose(centroid, t.Tensor(sp_centroid))
# Test complex with multiple ims
ims = t.rand((5,30,40)) + 1j * t.rand((5,30,40))
ims = t.rand((5, 30, 40)) + 1j * t.rand((5, 30, 40))
np_ims = ims.numpy()
sp_centroids = [ndimage.center_of_mass(np.abs(im)**2)
for im in np_ims]
centroids = image_processing.centroid_sq(ims, comp=True)
assert t.allclose(centroids, t.Tensor(np.array(sp_centroids)))
def test_sinc_subpixel_shift():
im = np.zeros((512,512), dtype=np.complex128)
im[256,256] = 1
im = np.zeros((512, 512), dtype=np.complex128)
im[256, 256] = 1
# test it by creating a single pixel object and seeing that it is
# shifted correctly
xs = np.arange(512) - 256
Ys,Xs = np.meshgrid(xs,xs)
sinc_im = np.sinc(Xs-0.3) * np.sinc(Ys-0.6)
Ys, Xs = np.meshgrid(xs, xs)
sinc_im = np.sinc(Xs - 0.3) * np.sinc(Ys - 0.6)
torch_im = t.as_tensor(im)
test_im = image_processing.sinc_subpixel_shift(torch_im,(0.3,0.6))
test_im = image_processing.sinc_subpixel_shift(torch_im, (0.3, 0.6))
# The fidelity isn't great due to the FFT-based approach, so we need
# a pretty relaxed condition
@@ -57,84 +58,82 @@ def test_sinc_subpixel_shift():
def test_find_pixel_shift():
# Test two real ims
big_im = t.rand((30,70))
im1 = big_im[3:,:-20]
im2 = big_im[:-3,20:]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([-3,20]))
big_im = t.rand((30, 70))
im1 = big_im[3:, :-20]
im2 = big_im[:-3, 20:]
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([-3, 20]))
# Test a real and complex im
big_im = t.rand((30,70))
im1 = big_im[:-5,10:].to(dtype=t.complex64)
im2 = big_im[5:,:-10]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,-10]))
assert t.all(image_processing.find_pixel_shift(im2,im1) == t.LongTensor([-5,10]))
big_im = t.rand((30, 70))
im1 = big_im[:-5, 10:].to(dtype=t.complex64)
im2 = big_im[5:, :-10]
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([5, -10]))
assert t.all(image_processing.find_pixel_shift(im2, im1) == t.LongTensor([-5, 10]))
# Test two complex ims
big_im = t.rand((45,45)) + 1j * t.rand((45,45))
im1 = big_im[:-5,:-4]
im2 = big_im[5:,4:]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,4]))
big_im = t.rand((45, 45)) + 1j * t.rand((45, 45))
im1 = big_im[:-5, :-4]
im2 = big_im[5:, 4:]
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([5, 4]))
def test_find_subpixel_shift():
# We can do this by creating a test probe and a test object
test_probe = t.rand((70,70)) + 1j * t.rand((70,70))
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
test_probe = t.rand((70, 70)) + 1j * t.rand((70, 70))
test_obj = t.ones((300, 300)) + 1j * t.rand((300, 300))
shift = t.tensor((0.8, 0.75))
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift, multiple_modes=False)
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0,0), resolution=50)
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0, 0), resolution=50)
# tolerance of 0.03 on this measurement
assert t.all(t.abs(shift - retrieved_shift) < 0.03)
def test_find_shift():
# We can do this by creating a test probe and a test object
test_probe = t.rand((200,200)) + 1j * t.rand((200,200))
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
test_probe = t.rand((200, 200)) + 1j * t.rand((200, 200))
test_obj = t.ones((300, 300)) + 1j * t.rand((300, 300))
shift = t.tensor((0.8, 0.75))
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift,
multiple_modes=False)[:-40,:-6]
multiple_modes=False)[:-40, :-6]
retrieved_shift = image_processing.find_shift(im, test_probe[40:,6:], resolution=50)
retrieved_shift = image_processing.find_shift(im, test_probe[40:, 6:], resolution=50)
# tolerance of 0.03 on this measurement
assert t.all(t.abs(shift + t.Tensor((40,6)) - retrieved_shift) < 0.03)
assert t.all(t.abs(shift + t.Tensor((40, 6)) - retrieved_shift) < 0.03)
def test_convolve_1d():
test_image = np.random.rand(400,300)
#test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
xs = np.linspace(-100,100,300)
kernel = 1/(1+xs**2)
test_image = np.random.rand(400, 300)
# test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
xs = np.linspace(-100, 100, 300)
kernel = 1 / (1 + xs**2)
# First, we test with everything real, dim=1
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
t.as_tensor(kernel),dim=1)
t.as_tensor(kernel), dim=1)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
assert np.allclose(convolved.numpy(),np_result)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image, axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
assert np.allclose(convolved.numpy(), np_result)
xs = np.linspace(-100,100,400)
kernel = 1/(1+xs**2)
xs = np.linspace(-100, 100, 400)
kernel = 1 / (1 + xs**2)
# Then with dim=0, and a non-fftshifted kernel
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
t.as_tensor(np.fft.ifftshift(kernel)),
fftshift_kernel=False)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0))
assert np.allclose(convolved.numpy(),np_result)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image, axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:, None], axis=0))
assert np.allclose(convolved.numpy(), np_result)
# And finally with complex input
convolved = image_processing.convolve_1d(t.as_tensor(test_image,dtype=t.complex64),
t.as_tensor(kernel,dtype=t.complex64)).numpy()
np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)
assert np.allclose(convolved,np_result)
convolved = image_processing.convolve_1d(t.as_tensor(test_image, dtype=t.complex64),
t.as_tensor(kernel, dtype=t.complex64)).numpy()
np_result = np.fft.ifft(np.fft.fft(test_image, axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:, None], axis=0)
assert np.allclose(convolved, np_result)
tests/tools/test_initializers.py
+58 -66
View File
@@ -1,26 +1,27 @@
from cdtools.tools import initializers
from cdtools.datasets import Ptycho2DDataset
import numpy as np
import torch as t
from cdtools.tools import initializers
from cdtools.datasets import Ptycho2DDataset
def test_exit_wave_geometry():
# First test a simple case where nothing need change
basis = t.Tensor([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(0,1)
shape = t.Size([73,56])
basis = t.Tensor([[0, -30e-6, 0],
[-20e-6, 0, 0]]).transpose(0, 1)
shape = t.Size([73, 56])
wavelength = 1e-9
distance = 1.
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
assert t.allclose(rs_basis[0,1],t.Tensor([-8.928571428571428e-07]))
assert t.allclose(rs_basis[1,0],t.Tensor([-4.5662100456621004e-07]))
assert t.allclose(rs_basis[0, 1], t.Tensor([-8.928571428571428e-07]))
assert t.allclose(rs_basis[1, 0], t.Tensor([-4.5662100456621004e-07]))
def test_calc_object_setup():
# First just try a simple case
probe_shape = t.Size([120,57])
translations = t.rand((30,2)) * 300
probe_shape = t.Size([120, 57])
translations = t.rand((30, 2)) * 300
t_max = t.max(translations, dim=0)[0]
t_min = t.min(translations, dim=0)[0]
obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations)
@@ -28,65 +29,58 @@ def test_calc_object_setup():
assert t.allclose(min_translation, t_min)
assert obj_shape == t.Size(exp_shape)
# Then add some padding
padding = 5
obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations, padding=padding)
assert t.allclose(min_translation, t_min - padding)
assert obj_shape == t.Size(exp_shape + 2 * padding)
def test_gaussian():
# Generate gaussian as a numpy array (square array)
shape = [10, 10]
sigma = [2.5, 2.5]
center = ((shape[0]-1)/2, (shape[1]-1)/2)
center = ((shape[0] - 1) / 2, (shape[1] - 1) / 2)
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
np_result = 10 * np.exp(-0.5 * ((x - center[1]) / sigma[1])**2 - 0.5 * ((y - center[0]) / sigma[0])**2)
init_result = initializers.gaussian(shape, sigma, amplitude=10).numpy()
assert np.allclose(init_result, np_result)
# Generate gaussian as a numpy array (rectangular array)
shape = [10, 5]
sigma = [2.5, 3]
center = ((shape[0]-1)/2, (shape[1]-1)/2)
center = ((shape[0] - 1) / 2, (shape[1] - 1) / 2)
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
np_result = np.exp(-0.5 * ((x - center[1]) / sigma[1])**2
- 0.5 * ((y - center[0]) / sigma[0])**2)
init_result = initializers.gaussian(shape, sigma).numpy()
assert np.allclose(init_result, np_result)
# Generate gaussian with curvature
shape = [20, 30]
sigma = [2.5, 5]
curvature = [1,0.6]
center = ((shape[0]-1)/2 + 3, (shape[1]-1)/2 - 1.4)
curvature = [1, 0.6]
center = ((shape[0] - 1) / 2 + 3, (shape[1] - 1) / 2 - 1.4)
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = (10+0j)*np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
np_result *= np.exp(0.5j*curvature[1]*(x-center[1])**2
+0.5j*curvature[0]*(y-center[0])**2)
init_result = initializers.gaussian(shape, sigma, center=center,
curvature=curvature, amplitude=10).numpy()
np_result = (10 + 0j) * np.exp(-0.5 * ((x - center[1]) / sigma[1])**2 - 0.5 * ((y - center[0]) / sigma[0])**2)
np_result *= np.exp(0.5j * curvature[1] * (x - center[1])**2 + 0.5j * curvature[0] * (y - center[0])**2)
init_result = initializers.gaussian(shape, sigma, center=center, curvature=curvature, amplitude=10).numpy()
assert np.allclose(init_result, np_result)
def test_gaussian_probe(ptycho_cxi_1):
dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
det_basis = t.Tensor(dataset.detector_geometry['basis'])
det_shape = t.Size(dataset.patterns.shape[-2:])
wavelength = dataset.wavelength
distance = dataset.detector_geometry['distance']
basis = initializers.exit_wave_geometry(det_basis,
det_shape,
wavelength,
distance)
det_shape,
wavelength,
distance)
# Basis is around 60nm in the i(y) direction, 85nm in the j(x) direction
# Full window is therefore about 15 um in i(y) and 20 um in the j(x) dir
@@ -94,15 +88,13 @@ def test_gaussian_probe(ptycho_cxi_1):
sigma = 5e-7
# Build a stage explicitly with numpy to compare against
x = (np.arange(256) - 127.5) * (-basis[0,1]).numpy()
y = (np.arange(256) - 127.5) * (-basis[1,0]).numpy()
Xs,Ys = np.meshgrid(x,y)
Rs = np.sqrt(Xs**2+Ys**2)
x = (np.arange(256) - 127.5) * (-basis[0, 1]).numpy()
y = (np.arange(256) - 127.5) * (-basis[1, 0]).numpy()
Xs, Ys = np.meshgrid(x, y)
Rs = np.sqrt(Xs**2 + Ys**2)
# Now we first test the non-propagated probe
np_probe = np.exp(-1/(2*sigma**2) * Rs**2)
np_probe = np.exp(- 1 / (2 * sigma**2) * Rs**2)
normalization = 0
for params, im in dataset:
@@ -110,27 +102,26 @@ def test_gaussian_probe(ptycho_cxi_1):
normalization /= len(dataset)
normalization_1 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
probe = initializers.gaussian_probe(
dataset, basis, det_shape, sigma).numpy()
assert np.allclose(probe, normalization_1*np_probe)
assert np.allclose(probe, normalization_1 * np_probe)
# And then a propagated probe
z = 1e-4 #nm
z = 1e-4 # nm
k = 2 * np.pi / wavelength
w0 = np.sqrt(2)*sigma
w0 = np.sqrt(2) * sigma
zr = np.pi * w0**2 / wavelength
wz = w0 * np.sqrt(1 + (z / zr)**2)
Rz = z * (1 + (zr / z)**2)
np_probe = np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
Rz = z * (1 + (zr / z)**2)
np_probe = np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
normalization_2 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
normalization_2 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
probe = initializers.gaussian_probe(dataset, basis, det_shape, sigma,
propagation_distance=z).numpy()
assert np.allclose(probe, normalization_2*np_probe)
assert np.allclose(probe, normalization_2 * np_probe)
def test_SHARP_style_probe(ptycho_cxi_1):
@@ -148,11 +139,13 @@ def test_SHARP_style_probe(ptycho_cxi_1):
wavelength,
distance)
assert basis.shape == t.Size([3, 2])
probe = initializers.SHARP_style_probe(dataset)
assert probe.shape == t.Size([256,256])
assert probe.shape == t.Size([256, 256])
probe = initializers.SHARP_style_probe(dataset, propagation_distance=20e-6)
assert probe.shape == t.Size([256,256])
assert probe.shape == t.Size([256, 256])
def test_RPI_spectral_init():
@@ -160,28 +153,27 @@ def test_RPI_spectral_init():
# since the original implementation is in numpy and there aren't any clear
# cases that can be calculated analytically.
pattern = np.random.rand(230,253).astype(np.float32)
probe = np.random.rand(230,253).astype(np.complex64)
obj_shape = [37,53]
pattern = np.random.rand(230, 253).astype(np.float32)
probe = np.random.rand(230, 253).astype(np.complex64)
obj_shape = [37, 53]
mask = t.Tensor(np.random.rand(*pattern.shape) > 0.04)
background = t.as_tensor(np.random.rand(*pattern.shape),dtype=t.float32) * 0.05
background = t.as_tensor(np.random.rand(*pattern.shape), dtype=t.float32) * 0.05
probe = t.as_tensor(probe)
pattern = t.as_tensor(pattern)
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape)
assert list(obj.shape) == [1]+obj_shape
assert list(obj.shape) == [1] + obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, mask=mask)
assert list(obj.shape) == [2]+obj_shape
assert list(obj.shape) == [2] + obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, background=background)
assert list(obj.shape) == [2]+obj_shape
assert list(obj.shape) == [2] + obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, mask=mask,
background=background)
assert list(obj.shape) == [2]+obj_shape
assert list(obj.shape) == [2] + obj_shape
tests/tools/test_interactions.py
+98 -114
View File
@@ -1,9 +1,10 @@
from cdtools.tools import interactions
import numpy as np
import torch as t
from numpy import fft
from numpy.fft import fftshift, ifftshift
from numpy import fft
import pytest
import torch as t
from cdtools.tools import interactions
# Have a random probe and a random object and test the two
@@ -14,59 +15,60 @@ import pytest
@pytest.fixture(scope='module')
def random_probe():
return np.random.rand(256,256) * np.exp(2j * np.pi * np.random.rand(256,256))
return np.random.rand(256, 256) * np.exp(2j * np.pi * np.random.rand(256, 256))
@pytest.fixture(scope='module')
def random_obj():
return np.random.rand(900,900) * np.exp(2j * np.pi * np.random.rand(900,900))
return np.random.rand(900, 900) * np.exp(2j * np.pi * np.random.rand(900, 900))
@pytest.fixture(scope='module')
def single_pixel_probe(scope='module'):
probe = np.zeros((256,256), dtype=np.complex128)
probe[128,128] = 1
probe = np.zeros((256, 256), dtype=np.complex128)
probe[128, 128] = 1
return probe
def test_translations_to_pixel():
# First, try the case where everything is ones and simple
basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
translations = t.rand((10,3))
basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t()
translations = t.rand((10, 3))
output = interactions.translations_to_pixel(basis, translations)
assert t.allclose(output, -translations[:,:2].flip(1))
assert t.allclose(output, -translations[:, :2].flip(1))
# Next, try a case with a single translation
translation = t.rand((3))
output = interactions.translations_to_pixel(basis, translation)
assert t.allclose(output, -translation[:2].flip(0))
# Then, try a case with no surface normal but with a real conversion
basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t()
translations = t.rand((10,3))
basis = t.Tensor([[0, -2, 0], [-1, 0, 0.1]]).t()
translations = t.rand((10, 3))
output = interactions.translations_to_pixel(basis, translations)
basis_vectors_inv = t.pinverse(basis)
translations[:,2] = 0 # manually project off z component
assert t.allclose(output, t.mm(translations,basis_vectors_inv.t()))
translations[:, 2] = 0 # manually project off z component
assert t.allclose(output, t.mm(translations, basis_vectors_inv.t()))
# Finally, try a case with a known surface normal (reflection)
basis = t.Tensor([[0,-1,0],[0,0,1]]).t()
surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)])
translations = t.rand((10,3))
basis = t.Tensor([[0, -1, 0], [0, 0, 1]]).t()
surface_normal = t.Tensor([np.sqrt(2), 0, -np.sqrt(2)])
translations = t.rand((10, 3))
output = interactions.translations_to_pixel(basis, translations,
surface_normal=surface_normal)
exp_translations = t.stack((-translations[:,1],translations[:,0]),dim=1)
exp_translations = t.stack((-translations[:, 1], translations[:, 0]), dim=1)
assert t.allclose(output, exp_translations)
def test_pixel_to_translations():
# First, try the case where everything is ones and simple
basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
translations = t.rand((10,3))
translations[:,2] = 0
basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t()
translations = t.rand((10, 3))
translations[:, 2] = 0
output = interactions.translations_to_pixel(basis, translations)
roundtrip = interactions.pixel_to_translations(basis, output)
assert t.allclose(translations, roundtrip)
# Next, try a case with a single translation
translation = t.rand((3))
translation[2] = 0
@@ -74,65 +76,59 @@ def test_pixel_to_translations():
roundtrip = interactions.pixel_to_translations(basis, output)
assert t.allclose(translation, roundtrip)
# Then, try a case with no surface normal but with a real conversion
basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t()
translations = t.rand((10,3))
translations[:,2] = 0 # manually project off z component
basis = t.Tensor([[0, -2, 0], [-1, 0, 0.1]]).t()
translations = t.rand((10, 3))
translations[:, 2] = 0 # manually project off z component
output = interactions.translations_to_pixel(basis, translations)
roundtrip = interactions.pixel_to_translations(basis, output)
assert t.allclose(translations, roundtrip)
# Finally, try a case with a known surface normal (reflection)
basis = t.Tensor([[0,-1,0],[0,0,1]]).t()
surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)])
translations = t.rand((10,3))
translations[:,2] = 0 # manually project off z component
basis = t.Tensor([[0, -1, 0], [0, 0, 1]]).t()
surface_normal = t.Tensor([np.sqrt(2), 0, -np.sqrt(2)])
translations = t.rand((10, 3))
translations[:, 2] = 0 # manually project off z component
output = interactions.translations_to_pixel(basis, translations,
surface_normal=surface_normal)
roundtrip = interactions.pixel_to_translations(basis, output,
surface_normal=surface_normal)
surface_normal=surface_normal)
assert t.allclose(translations, roundtrip)
def test_project_translations_to_sample():
# First, try the case where everything is ones and simple
basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
translations = t.rand((10,3))
basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t()
translations = t.rand((10, 3))
pixels, props = interactions.project_translations_to_sample(basis, translations)
assert np.allclose(pixels[:,0].numpy(),-translations[:,1])
assert np.allclose(pixels[:,1].numpy(),-translations[:,0])
assert np.allclose(props.numpy(),-translations[:,2:].numpy())
assert np.allclose(pixels[:, 0].numpy(), -translations[:, 1])
assert np.allclose(pixels[:, 1].numpy(), -translations[:, 0])
assert np.allclose(props.numpy(), -translations[:, 2:].numpy())
# Next, a simple tilt along one axis. This is a 45 degree rotation
# around the positive y-axis
# Thus, y-axis translations are unaffected, but x-axis translations
# induce a motion of 1/sqrt(2) in the j- pixel space, as well as
# creating a propagation (negative propagation for positive x)
basis = t.Tensor([[0,-1e-3,0],[-np.sqrt(2)*1e-3,0,np.sqrt(2)*1e-3]]).t()
translations = t.rand((10,3))
basis = t.Tensor([[0, -1e-3, 0], [-np.sqrt(2) * 1e-3, 0, np.sqrt(2) * 1e-3]]).t()
translations = t.rand((10, 3))
pixels, props = interactions.project_translations_to_sample(basis, translations)
print(props.numpy())
print(-translations[:,2:].numpy() - translations[:,:1].numpy())
assert np.allclose(pixels[:,0].numpy(),-translations[:,1]*1e3)
assert np.allclose(pixels[:,1].numpy(),-translations[:,0]*1e3/np.sqrt(2))
assert np.allclose(props.numpy(),-translations[:,2:].numpy() - translations[:,:1].numpy())
print(-translations[:, 2:].numpy() - translations[:, :1].numpy())
assert np.allclose(pixels[:, 0].numpy(), -translations[:, 1] * 1e3)
assert np.allclose(pixels[:, 1].numpy(), -translations[:, 0] * 1e3 / np.sqrt(2))
assert np.allclose(props.numpy(), -translations[:, 2:].numpy() - translations[:, :1].numpy())
# Finally, we check a non-orthogonal case
def test_ptycho_2D_round(random_probe, random_obj):
# Test a stack of images
translations = np.random.rand(10,2) * 500
exit_waves_np = [random_probe * \
random_obj[tr[0]:tr[0]+random_probe.shape[0],
tr[1]:tr[1]+random_probe.shape[1]] for
translations = np.random.rand(10, 2) * 500
exit_waves_np = [random_probe * random_obj[tr[0]:tr[0] + random_probe.shape[0],
tr[1]:tr[1] + random_probe.shape[1]] for
tr in np.round(translations).astype(int)]
exit_waves_t = interactions.ptycho_2D_round(t.as_tensor(random_probe),
t.as_tensor(random_obj),
@@ -146,14 +142,13 @@ def test_ptycho_2D_round(random_probe, random_obj):
assert np.allclose(exit_wave_t.numpy(), exit_waves_np[0])
def test_ptycho_2D_linear(single_pixel_probe, random_obj):
# For this one, I just want to check one translation, but
# I need to check both formats
translations = np.array([[46.7,53.2]])
translation = np.array([46.7,53.2])
translations = np.array([[46.7, 53.2]])
translation = np.array([46.7, 53.2])
exit_waves_probe = interactions.ptycho_2D_linear(
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
@@ -167,14 +162,9 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
shift_probe=True)
# Check that the outputs match
assert t.allclose(exit_waves_probe[0],exit_wave_probe)
assert t.allclose(exit_waves_probe[0], exit_wave_probe)
exit_waves_obj = interactions.ptycho_2D_linear(
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.tensor(translations),
shift_probe=False)
exit_waves_obj = interactions.ptycho_2D_linear(t.as_tensor(single_pixel_probe), t.as_tensor(random_obj), t.tensor(translations), shift_probe=False)
exit_wave_obj = interactions.ptycho_2D_linear(
t.as_tensor(single_pixel_probe),
@@ -183,38 +173,37 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
shift_probe=False)
# Check that the outputs match
assert t.allclose(exit_waves_obj[0],exit_wave_obj)
assert t.allclose(exit_waves_obj[0], exit_wave_obj)
# For the shifted probe, we should find 4 pixels with intensity
exit_waves_probe = t.as_tensor(exit_waves_probe)[0]
probe_shift = np.array([[0.3*0.8,0.3*0.2],
[0.7*0.8,0.7*0.2]])
obj_section = random_obj[128+46:128+48,
128+53:128+55]
exit_section = exit_waves_probe[128:130,128:130]
probe_shift = np.array([[0.3 * 0.8, 0.3 * 0.2],
[0.7 * 0.8, 0.7 * 0.2]])
obj_section = random_obj[128 + 46:128 + 48,
128 + 53:128 + 55]
exit_section = exit_waves_probe[128:130, 128:130]
assert np.allclose(probe_shift * obj_section, exit_section)
# For the shifted obj, we should find one pixel with intensity
exit_waves_obj = t.as_tensor(exit_waves_obj)[0]
obj_shift = np.array([[0.3*0.8,0.3*0.2],
[0.7*0.8,0.7*0.2]])
obj_section = random_obj[128+46:128+48,
128+53:128+55]
exit_pixel = exit_waves_obj[128,128]
assert np.isclose(np.sum(obj_shift * obj_section),exit_pixel)
obj_shift = np.array([[0.3 * 0.8, 0.3 * 0.2],
[0.7 * 0.8, 0.7 * 0.2]])
obj_section = random_obj[128 + 46:128 + 48,
128 + 53:128 + 55]
exit_pixel = exit_waves_obj[128, 128]
assert np.isclose(np.sum(obj_shift * obj_section), exit_pixel)
# Test for a single translation
def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
# For this one, I just want to check one translation, but
# I need to check both formats
translations = np.array([[46.7,53.2]])
translation = np.array([46.7,53.2])
translations = np.array([[46.7, 53.2]])
translation = np.array([46.7, 53.2])
exit_waves_probe = interactions.ptycho_2D_sinc(
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
@@ -228,32 +217,31 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
shift_probe=True)
# Check that the outputs match
assert t.allclose(exit_waves_probe[0],exit_wave_probe)
assert t.allclose(exit_waves_probe[0], exit_wave_probe)
# Now we explicitly define what the sinc interpolated array should
# look like
xs = np.arange(256) - 128
Ys,Xs = np.meshgrid(xs,xs)
Ys, Xs = np.meshgrid(xs, xs)
sinc_probe = np.sinc(Xs) * np.sinc(Ys)
# Just check that the unshifted probe is correct
assert np.allclose(single_pixel_probe, sinc_probe)
sinc_shifted_probe = np.sinc(Xs-0.7) * np.sinc(Ys-0.2)
obj_section = random_obj[46:46+256,
53:53+256]
sinc_shifted_probe = np.sinc(Xs - 0.7) * np.sinc(Ys - 0.2)
obj_section = random_obj[46:46 + 256,
53:53 + 256]
exit_wave_np = sinc_shifted_probe * obj_section
exit_wave_torch = exit_wave_probe.numpy()
# The fidelity isn't great due to the FFT-based approach, so we need
# a pretty relaxed condition
assert np.max(np.abs(exit_wave_np-exit_wave_torch)) < 0.005
assert np.max(np.abs(exit_wave_np - exit_wave_torch)) < 0.005
def test_RPI_interaction(random_probe, random_obj):
random_obj1 = random_obj[:79,:68] * 0 + 1
random_obj1 = random_obj[:79, :68] * 0 + 1
random_probe1 = random_probe * 0 + 1
t_random_obj1 = t.as_tensor(random_obj1)
t_random_probe1 = t.as_tensor(random_probe1)
@@ -261,36 +249,32 @@ def test_RPI_interaction(random_probe, random_obj):
obj1_fourier = fftshift(fft.fft2(ifftshift(random_obj1), norm='ortho'))
obj1_ups = np.zeros(random_probe1.shape[:2]).astype(np.complex128)
obj1_ups[random_probe1.shape[0]//2 - 79//2:
-(random_probe1.shape[0]-79 - (random_probe1.shape[0]//2 - 79//2)),
(random_probe1.shape[1]-68)//2:
(random_probe1.shape[1]-68)//2 + 68] = obj1_fourier
obj1_ups[random_probe1.shape[0] // 2 - 79 // 2:
-(random_probe1.shape[0] - 79 - (random_probe1.shape[0] // 2 - 79 // 2)),
(random_probe1.shape[1] - 68) // 2:
(random_probe1.shape[1] - 68) // 2 + 68] = obj1_fourier
output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups),
norm='ortho'))
norm='ortho'))
output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1]
/ (random_obj1.shape[-2] * random_obj1.shape[-1]))
output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1] / (random_obj1.shape[-2] * random_obj1.shape[-1]))
assert np.allclose(t_output1, output1)
random_obj2 = np.stack([random_obj[:64,:89]]*3)
random_probe2 = random_probe[3:,5:]
random_obj2 = np.stack([random_obj[:64, :89]] * 3)
random_probe2 = random_probe[3:, 5:]
t_random_obj2 = t.as_tensor(random_obj2)
t_random_probe2 = t.as_tensor(random_probe2)
t_output2 = interactions.RPI_interaction(t_random_probe2, t_random_obj2)
obj2_fourier = fftshift(fft.fft2(ifftshift(random_obj2), norm='ortho'))
obj2_ups = np.zeros((3,)+random_probe2.shape[:2]).astype(np.complex128)
obj2_ups[:,(random_probe2.shape[0]-64)//2:
(random_probe2.shape[0]-64)//2 + 64,
(random_probe2.shape[1]-89)//2:
(random_probe2.shape[1]-89)//2 + 89] = obj2_fourier
obj2_ups = np.zeros((3,) + random_probe2.shape[:2]).astype(np.complex128)
obj2_ups[:, (random_probe2.shape[0] - 64) // 2:
(random_probe2.shape[0] - 64) // 2 + 64,
(random_probe2.shape[1] - 89) // 2:
(random_probe2.shape[1] - 89) // 2 + 89] = obj2_fourier
output2 = random_probe2 * fftshift(fft.ifft2(ifftshift(obj2_ups),
norm='ortho'))
norm='ortho'))
output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1]
/ (random_obj2.shape[-2] * random_obj2.shape[-1]))
output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1] / (random_obj2.shape[-2] * random_obj2.shape[-1]))
assert np.allclose(t_output2, output2)
tests/tools/test_losses.py
+32 -38
View File
@@ -1,79 +1,73 @@
from cdtools.tools import losses
import numpy as np
import torch as t
from cdtools.tools import losses
# The idea here is to use a simple numpy calculation of the various
# objective functions to check the torch implementations and make sure
# that any optimizations in the future don't change the results
def test_amplitude_mse():
# Make some fake data
data = np.random.rand(10,100,100)
data = np.random.rand(10, 100, 100)
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
sim = data + 0.1 * np.random.rand(10, 100, 100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(bool)
mask = (np.random.rand(100, 100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2)
#np_result /= data.size
torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
# np_result /= data.size
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2)
#np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim),
mask = t.from_numpy(mask))
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
# np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
def test_intensity_mse():
# Make some fake data
data = np.random.rand(10,100,100)
data = np.random.rand(10, 100, 100)
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
sim = data + 0.1 * np.random.rand(10, 100, 100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(bool)
mask = (np.random.rand(100, 100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum((data - sim)**2)
np_result /= data.size
torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
np_result /= data.size
torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (data - sim)**2)
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim),
mask = t.from_numpy(mask))
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
def test_poisson_nll():
# Make some fake data
data = np.random.rand(10,100,100)
data = np.random.rand(10, 100, 100)
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
sim = data + 0.1 * np.random.rand(10, 100, 100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(bool)
mask = (np.random.rand(100, 100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum(sim - data * np.log(sim))
np_result /= data.size
torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
np_result /= data.size
torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (sim - data * np.log(sim)))
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim),
mask = t.from_numpy(mask), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim),
mask=t.from_numpy(mask), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
tests/tools/test_measurements.py
+37 -46
View File
@@ -1,99 +1,90 @@
from cdtools.tools import measurements
import torch as t
import numpy as np
from cdtools.tools import measurements
def test_intensity():
wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
epsilon=1e-6
wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10))
epsilon = 1e-6
np_result = np.abs(wavefields.numpy())**2 + epsilon
assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon),
assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon),
t.as_tensor(np_result))
# Test single field case
assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon),
assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon),
t.as_tensor(np_result[0]))
det_slice = np.s_[3:,5:8]
assert t.allclose(measurements.intensity(wavefields,det_slice,epsilon=epsilon),
t.as_tensor(np_result[(np.s_[:],)+det_slice]))
det_slice = np.s_[3:, 5:8]
assert t.allclose(measurements.intensity(wavefields, det_slice, epsilon=epsilon),
t.as_tensor(np_result[(np.s_[:],) + det_slice]))
# Test single field case
assert t.allclose(measurements.intensity(wavefields[0],det_slice,epsilon=epsilon),
assert t.allclose(measurements.intensity(wavefields[0], det_slice, epsilon=epsilon),
t.as_tensor(np_result[0][det_slice]))
# With oversampling on
np_oversampling_result = (np_result[:,::2,::2] + \
np_result[:,1::2,::2] + \
np_result[:,::2,1::2] + \
np_result[:,1::2,1::2]) / 4
np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4
# With multiple fields
assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon, oversampling=2),
assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result,))
# With a single field
assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon, oversampling=2),
assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result[0],))
def test_incoherent_sum():
# With no explicit slice given
wavefields = t.rand((5,4,10,10)) + 1j * t.rand((5,4,10,10))
epsilon=1e-6
np_result = np.sum(np.abs(wavefields.numpy())**2,axis=-3) + epsilon
assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon),
wavefields = t.rand((5, 4, 10, 10)) + 1j * t.rand((5, 4, 10, 10))
epsilon = 1e-6
np_result = np.sum(np.abs(wavefields.numpy())**2, axis=-3) + epsilon
assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon),
t.as_tensor(np_result))
# Test single field case
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon),
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon),
t.as_tensor(np_result[0]))
# With a slice given
det_slice = np.s_[3:,5:8]
assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon),
t.as_tensor(np_result[(np.s_[:],)+det_slice]))
det_slice = np.s_[3:, 5:8]
assert t.allclose(measurements.incoherent_sum(wavefields, det_slice, epsilon=epsilon),
t.as_tensor(np_result[(np.s_[:],) + det_slice]))
# Test single field case
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],det_slice,epsilon=epsilon),
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], det_slice, epsilon=epsilon),
t.as_tensor(np_result[0][det_slice]))
# With oversampling on
np_oversampling_result = (np_result[:,::2,::2] + \
np_result[:,1::2,::2] + \
np_result[:,::2,1::2] + \
np_result[:,1::2,1::2]) / 4
np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4
# With multiple fields
assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon, oversampling=2),
assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result,))
# With a single field
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon, oversampling=2),
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result[0],))
def test_quadratic_background():
# test with intensity
wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
epsilon=1e-6
background = t.rand((10,10))
wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10))
epsilon = 1e-6
background = t.rand((10, 10))
np_result = np.abs(wavefields.numpy())**2 + background.numpy()**2 + epsilon
det_slice = np.s_[3:,5:8]
det_slice = np.s_[3:, 5:8]
result = measurements.quadratic_background(wavefields,background[det_slice],
result = measurements.quadratic_background(wavefields, background[det_slice],
detector_slice=det_slice,
epsilon=epsilon,
measurement=measurements.intensity)
assert t.allclose(result, t.tensor(np_result[(np.s_[:],)+det_slice]))
assert t.allclose(result, t.tensor(np_result[(np.s_[:],) + det_slice]))
# test with incoherent sum but no slice and no stack
wavefields = t.rand((4,10,10)) + 1j * t.rand((4,10,10))
np_result = np.sum(np.abs(wavefields.numpy())**2,axis=0)
wavefields = t.rand((4, 10, 10)) + 1j * t.rand((4, 10, 10))
np_result = np.sum(np.abs(wavefields.numpy())**2, axis=0)
np_result += background.numpy()**2
result = measurements.quadratic_background(wavefields, background,
epsilon=epsilon,
tests/tools/test_plotting.py
+15 -12
View File
@@ -1,47 +1,50 @@
from cdtools.tools import plotting
from cdtools.tools import initializers
import numpy as np
import torch as t
import scipy.datasets
import matplotlib.pyplot as plt
from cdtools.tools import plotting
from cdtools.tools import initializers
def test_plot_amplitude(show_plot):
# Test with tensor
im = t.as_tensor(scipy.datasets.ascent(),dtype=t.complex128)
plotting.plot_amplitude(im, basis = np.array([[0,-1], [-1,0], [0,0]]), title = 'Test Amplitude')
im = t.as_tensor(scipy.datasets.ascent(), dtype=t.complex128)
plotting.plot_amplitude(im, basis=np.array([[0, -1], [-1, 0], [0, 0]]), title='Test Amplitude')
if show_plot:
plt.show()
# Test with numpy array
im = scipy.datasets.ascent().astype(np.complex128)
plotting.plot_amplitude(im, title = 'Test Amplitude')
plotting.plot_amplitude(im, title='Test Amplitude')
if show_plot:
plt.show()
def test_plot_phase(show_plot):
# Test with tensor
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
plotting.plot_phase(im, title = 'Test Phase')
im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
plotting.plot_phase(im, title='Test Phase')
if show_plot:
plt.show()
# Test with numpy array
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]).numpy()
plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[0,-1], [-1,0], [0,0]]))
im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1]).numpy()
plotting.plot_phase(im, title='Test Phase', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
if show_plot:
plt.show()
def test_plot_colorized(show_plot):
# Test with tensor
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
gaussian = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
im = gaussian * t.as_tensor(scipy.datasets.ascent(), dtype=t.complex64)
plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[0,-1], [-1,0], [0,0]]))
plotting.plot_colorized(im, title='Test Colorize', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
if show_plot:
plt.show()
# Test with numpy array
im = im.numpy()
plotting.plot_colorized(im, title = 'Test Colorize')
plotting.plot_colorized(im, title='Test Colorize')
if show_plot:
plt.show()
tests/tools/test_propagators.py
+194 -211
View File
@@ -1,7 +1,3 @@
from cdtools.tools import initializers
from cdtools.tools import propagators
from cdtools.tools import image_processing
import numpy as np
import torch as t
import pytest
@@ -9,35 +5,37 @@ import scipy.datasets
from scipy import stats
from matplotlib import pyplot as plt
from cdtools.tools import initializers
from cdtools.tools import propagators
from cdtools.tools import image_processing
@pytest.fixture(scope='module')
def exit_waves_1():
# Import scipy test image and add a random phase
obj = scipy.datasets.ascent()[0:64,0:64].astype(np.complex128)
arr = np.random.random_sample((64,64))
obj *= (arr+(1-arr**2)**.5*1j)
obj = scipy.datasets.ascent()[0:64, 0:64].astype(np.complex128)
arr = np.random.random_sample((64, 64))
obj *= (arr + (1 - arr**2)**.5 * 1j)
obj = t.as_tensor(obj)
# Construct wavefront from image
probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3)
return probe * obj
def test_far_field(exit_waves_1):
# Far field diffraction patterns calculated by numpy with zero frequency in center
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho'))
assert(np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()), norm='ortho'))
assert (np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
def test_inverse_far_field(exit_waves_1):
# We want the inverse far field to map back to the exit waves with no intensity corrections
# Far field result for exit waves calculated with numpy
far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho')))
far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()), norm='ortho')))
assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
assert (np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
def test_generate_high_NA_k_intensity_map():
@@ -45,9 +43,9 @@ def test_generate_high_NA_k_intensity_map():
# We need to generate a plausible scenario. I will start
# by using the initializer to generate a reasonable exit wave geometry
# and detector pair
basis = t.Tensor([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(0,1)
shape = t.Size([478,573])
basis = t.Tensor([[0, -30e-6, 0],
[-20e-6, 0, 0]]).transpose(0, 1)
shape = t.Size([478, 573])
wavelength = 1e-9
distance = 1
rs_basis = \
@@ -60,7 +58,7 @@ def test_generate_high_NA_k_intensity_map():
# generate a good test exit wave
i = (np.arange(478) - 240)
j = (np.arange(573) - 270)
Is,Js = np.meshgrid(i,j,indexing='ij')
Is, Js = np.meshgrid(i, j, indexing='ij')
wavefield = ((np.abs(Is) < 20) * (np.abs(Js) < 25)).astype(np.complex128)
t_wavefield = t.as_tensor(wavefield, dtype=t.complex64)
@@ -74,37 +72,36 @@ def test_generate_high_NA_k_intensity_map():
# Checking first that for a low-NA propagation they give the same result
# 1e-4 tolerance seems to be reasonable in this comparison given my
# exploration with the code
#assert np.max(np.abs(high_NA-low_NA))/np.max(np.abs(low_NA)) < 1e-4
# assert np.max(np.abs(high_NA-low_NA))/np.max(np.abs(low_NA)) < 1e-4
# Now I will explore some results with a tilted sample
#print(rs_basis)
#print(rs_basis_tilted)
distance = 0.01#6e-3
rs_basis = \
initializers.exit_wave_geometry(basis, shape, wavelength, distance)
# print(rs_basis)
# print(rs_basis_tilted)
distance = 0.01 # 6e-3
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
rs_basis_tilted = rs_basis.clone()
rs_basis_tilted[2,1] = rs_basis_tilted[0,1]
rs_basis_tilted[2, 1] = rs_basis_tilted[0, 1]
k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
rs_basis_tilted, basis, shape, distance, wavelength,
dtype=t.float32)
high_NA_propagated = propagators.high_NA_far_field(
t_wavefield, k_map, intensity_map=intensity_map)
low_NA_propagated = propagators.far_field(t_wavefield)
low_NA = low_NA_propagated.numpy()
high_NA = high_NA_propagated.numpy()
#plt.close('all')
#plt.imshow(np.abs(low_NA))
#plt.figure()
#plt.imshow(np.abs(high_NA))
#plt.colorbar()
#plt.imshow(np.abs(wavefield))
#plt.show()
print(low_NA.shape, high_NA.shape)
# plt.close('all')
# plt.imshow(np.abs(low_NA))
# plt.figure()
# plt.imshow(np.abs(high_NA))
# plt.colorbar()
# plt.imshow(np.abs(wavefield))
# plt.show()
# Now I want to test that it doesn't crash for wavefields of various shapes
propagators.high_NA_far_field(t_wavefield.unsqueeze(0),
@@ -114,8 +111,9 @@ def test_generate_high_NA_k_intensity_map():
# I believe this works, but I still would like to get a second method for
# simulating at least one diffraction pattern as an independent check
#assert 0
# assert 0
def test_near_field_direction():
#
@@ -135,19 +133,19 @@ def test_near_field_direction():
x = (np.arange(901) - 400)
y = (np.arange(1200) - 500)
Ys,Xs = np.meshgrid(y,x)
Rs = np.sqrt(Xs**2+Ys**2)
E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)),dtype=t.complex64)
Ys, Xs = np.meshgrid(y, x)
Rs = np.sqrt(Xs**2 + Ys**2)
E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)), dtype=t.complex64)
E0_real = propagators.inverse_far_field(E0_fourier)
# This is in the top-left corner in Fourier space
wavelength = 3e-9 #nm
z = 1000e-9
wavelength = 3e-9 # nm
z = 1000e-9 # nm
asp = propagators.generate_angular_spectrum_propagator(
E0_real.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex64)
E0_real.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex64)
Ez_real = propagators.near_field(E0_real,asp)
Ez_real = propagators.near_field(E0_real, asp)
centroid = image_processing.centroid(t.abs(Ez_real))
@@ -155,35 +153,35 @@ def test_near_field_direction():
assert centroid[0] < Ez_real.shape[0] // 2
# Assert it's in left half
assert centroid[1] < Ez_real.shape[1] // 2
#plt.imshow(t.abs(E0_fourier))
#plt.figure()
#plt.imshow(t.abs(E0_real))
#plt.figure()
#plt.imshow(t.abs(Ez_real))
#plt.show()
# plt.imshow(t.abs(E0_fourier))
# plt.figure()
# plt.imshow(t.abs(E0_real))
# plt.figure()
# plt.imshow(t.abs(Ez_real))
# plt.show()
def test_near_field():
# The strategy is to compare the propagation of a gaussian beam to
# the propagation in the paraxial approximation.
x = (np.arange(901) - 450) * 1.5e-9
y = (np.arange(1200) - 600) * 1e-9
Ys,Xs = np.meshgrid(y,x)
Rs = np.sqrt(Xs**2+Ys**2)
Ys, Xs = np.meshgrid(y, x)
Rs = np.sqrt(Xs**2 + Ys**2)
wavelength = 3e-9 #nm
sigma = 20e-9 #nm
z = 1000e-9 #nm
wavelength = 3e-9 # nm
sigma = 20e-9 # nm
z = 1000e-9 # nm
k = 2 * np.pi / wavelength
w0 = np.sqrt(2)*sigma
w0 = np.sqrt(2) * sigma
zr = np.pi * w0**2 / wavelength
wz = w0 * np.sqrt(1 + (z / zr)**2)
Rz = z * (1 + (zr / z)**2)
Rz = z * (1 + (zr / z)**2)
E0 = np.exp(-Rs**2 / w0**2)
# The analytical expression for propagation of a gaussian beam in the
@@ -208,52 +206,49 @@ def test_near_field():
# If we choose e^(-ikx) to represent light propagating along K, the
# answer is no, and we find we have to use the inverse FT instead.
# Thus, e^(ikx) is the right choice here.
Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * ( z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr))
Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * (z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr))
Ez_nozphase = Ez * np.exp(-1j * k * z)
# First we check it normally
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
assert np.max(np.abs(Ez_nozphase-Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
assert np.max(np.abs(Ez_nozphase - Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
Emz = np.conj(Ez_nozphase)
Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy()
Emz_t = propagators.inverse_near_field(t.as_tensor(E0), asp).numpy()
# Again, 10^-3 is about all the accuracy we can expect
assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz))
assert np.max(np.abs(Emz - Emz_t)) < 1e-3 * np.max(np.abs(Emz))
# Then, we check that the bandlimiting at least does something
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,
E0.shape, (1.5e-9, 1e-9), wavelength, z,
dtype=t.complex128, bandlimit=0.3)
assert asp[140,0] == 0
assert asp[0,180] == 0
assert asp[130,0] != 0
assert asp[0,175] != 0
assert asp[140, 0] == 0
assert asp[0, 180] == 0
assert asp[130, 0] != 0
assert asp[0, 175] != 0
# Then, we check that automatic differentiation works
z = t.tensor([z],requires_grad=True)
spacing = t.tensor((1.5e-9,1e-9), requires_grad=True)
z = t.tensor([z], requires_grad=True)
spacing = t.tensor((1.5e-9, 1e-9), requires_grad=True)
wavelength = t.tensor([wavelength], requires_grad=True)
asp = propagators.generate_angular_spectrum_propagator(
E0.shape, spacing, wavelength, z)
t.real(asp[10,10]).backward()
t.real(asp[10, 10]).backward()
assert z.grad != 0
assert spacing.grad[0] != 0
assert wavelength.grad !=0
assert wavelength.grad != 0
def test_generalized_near_field():
@@ -264,26 +259,25 @@ def test_generalized_near_field():
# First, we should do a test with the phase ramp along the z direction
# explicitly included
basis= np.array([[0,-1.5e-9],[-1e-9,0],[0,0]])
i_vec,j_vec = np.arange(901) - 450 ,np.arange(1200)-600
Is, Js = np.meshgrid(i_vec,j_vec,indexing='ij')
Xs_0,Ys_0,Zs_0 = np.tensordot(basis,np.stack([Is,Js]),axes=1)
#x = (np.arange(901) - 450) * 1.5e-9
#y = (np.arange(1200) - 600) * 1e-9
#Xs_0,Ys_0 = np.meshgrid(x,y)
#Zs_0 = np.zeros(Xs_0.shape)
basis = np.array([[0, -1.5e-9], [-1e-9, 0], [0, 0]])
i_vec, j_vec = np.arange(901) - 450, np.arange(1200) - 600
Is, Js = np.meshgrid(i_vec, j_vec, indexing='ij')
Xs_0, Ys_0, Zs_0 = np.tensordot(basis, np.stack([Is, Js]), axes=1)
# x = (np.arange(901) - 450) * 1.5e-9
# y = (np.arange(1200) - 600) * 1e-9
# Xs_0,Ys_0 = np.meshgrid(x,y)
# Zs_0 = np.zeros(Xs_0.shape)
Positions = np.stack([Xs_0, Ys_0, Zs_0])
Positions = np.stack([Xs_0,Ys_0,Zs_0])
# assert 0
wavelength = 3e-9 #nm
sigma = 20e-9 #nm
z = 1000e-9 #nm
wavelength = 3e-9 # nm
sigma = 20e-9 # nm
z = 1000e-9 # nm
k = 2 * np.pi / wavelength
w0 = np.sqrt(2)*sigma
w0 = np.sqrt(2) * sigma
zr = np.pi * w0**2 / wavelength
# The analytical expression for propagation of a gaussian beam in the
# paraxial approx
@@ -292,140 +286,132 @@ def test_generalized_near_field():
def get_inv_R(Zs):
return Zs / (Zs**2 + zr**2)
def get_E(Xs, Ys, Zs, correct=True):
# Again, this follows the convention opposite from Wikipedia. See
# the note in test_angular_spectrum_propagator.
# if correct is True, remove the e^(ikz) dependence
Rs_sq = Xs**2 + Ys**2
Wzs = get_w(Zs)
E = w0 / Wzs * np.exp(-Rs_sq / Wzs**2) *\
np.exp(1j * k * ( Zs + Rs_sq * get_inv_R(Zs) / 2) + \
- 1j * np.arctan(Zs / zr))
np.exp(1j * k * (Zs + Rs_sq * get_inv_R(Zs) / 2) - 1j * np.arctan(Zs / zr))
# This removes the z-dependence of the phase
if correct:
if correct:
E = E * np.exp(-1j * k * Zs)
return E
def check_equiv(analytical, numerical):
phase = np.angle(np.mean(numerical.conj()*analytical))
comp = np.exp(1j*phase) * numerical
return (np.max(np.abs(analytical-comp))
< 1e-3 * np.max(np.abs(analytical)))
phase = np.angle(np.mean(numerical.conj() * analytical))
comp = np.exp(1j * phase) * numerical
return (np.max(np.abs(analytical - comp)) < 1e-3 * np.max(np.abs(analytical)))
# We make some rotation matrices to test
# This tests the straight ahead case
I = np.eye(3)
IdentityMatrix = np.eye(3)
# This tests a rotation about the y axis
th = np.deg2rad(5)
Ry = np.array([[np.cos(th),0,np.sin(th)],
[0,1,0],
[-np.sin(th),0,np.cos(th)]])
Ry = np.array([[np.cos(th), 0, np.sin(th)],
[0, 1, 0],
[-np.sin(th), 0, np.cos(th)]])
# This tests a rotation about two axes
phi = np.deg2rad(2)
Rx = np.array([[1,0,0],
[0,np.cos(phi),-np.sin(phi)],
[0,np.sin(phi),np.cos(phi)]])
Rboth = np.matmul(Rx,Ry)
Rx = np.array([[1, 0, 0],
[0, np.cos(phi), -np.sin(phi)],
[0, np.sin(phi), np.cos(phi)]])
Rboth = np.matmul(Rx, Ry)
# This tests a shearing
shear = 0.23
Rshear = np.array([[1,shear,0],
[0,1,0],
[0,0,1]])
Rshear = np.array([[1, shear, 0],
[0, 1, 0],
[0, 0, 1]])
# This tests an inversion of the axes
Rinv = np.array([[-1,0,0],
[0,-1,0],
[0,0,-1]])
Rinv = np.array([[-1, 0, 0],
[0, -1, 0],
[0, 0, -1]])
# This tests a reflection about the y-z plane
Rrefl = np.array([[-1,0,0],
[0,1,0],
[0,0,-1]])
# This tests a shearing and a rotation together
Rall = np.matmul(Rrefl,np.matmul(Rboth,Rshear))
Rrefl = np.array([[-1, 0, 0],
[0, 1, 0],
[0, 0, -1]])
# This tests a shearing and a rotation together
Rall = np.matmul(Rrefl, np.matmul(Rboth, Rshear))
# And we make some propagation vectors to test:
# This is along the z direction
z_dir = np.array([0,0,1])
z_dir = np.array([0, 0, 1])
# This checks that it's not sensitive to the magnitude
z_dir_large = np.array([0,0,10])
z_dir_large = np.array([0, 0, 10])
# And finally some offset vectors
# This checks straight ahead
z_offset = np.array([0,0,z])
z_offset = np.array([0, 0, z])
# This checks with an offset in x and y
shear_offset = np.array([0.1*z,-0.03*z,z])
shear_offset = np.array([0.1 * z, -0.03 * z, z])
# This checks with an offset in x and y, with negative z
shear_back_offset = np.array([0.1*z,-0.03*z,-z])
shear_back_offset = np.array([0.1 * z, -0.03 * z, -z])
rot_mats = [Rrefl,I,Rinv, Rboth, Rboth,Rboth, Rall, Rall, Rall, I, Rall]
offset_vecs = [z_offset]*8 + [shear_offset] + [shear_back_offset]*2
propagation_vecs = ['perp','offset',z_dir,
'perp','offset',z_dir_large,
'perp','offset',z_dir_large,
rot_mats = [Rrefl, IdentityMatrix, Rinv, Rboth, Rboth, Rboth, Rall, Rall, Rall, IdentityMatrix, Rall]
offset_vecs = [z_offset] * 8 + [shear_offset] + [shear_back_offset] * 2
propagation_vecs = ['perp', 'offset', z_dir,
'perp', 'offset', z_dir_large,
'perp', 'offset', z_dir_large,
z_dir, z_dir_large]
purposes = ['standard']*3 + ['both-rot']*3 + ['shear-rot']*3 + ['backward']*2
#rot_mats = [Ry.transpose()]
#offset = np.cross(np.dot(Ry.transpose(),basis)[:,0],
purposes = ['standard'] * 3 + ['both-rot'] * 3 + ['shear-rot'] * 3 + ['backward'] * 2
# rot_mats = [Ry.transpose()]
# offset = np.cross(np.dot(Ry.transpose(),basis)[:,0],
# np.dot(Ry.transpose(),basis)[:,1])
#offset /= np.linalg.norm(offset) / 3e-6
#offset_vecs = [-offset]#[shear_offset]
#propagation_vecs = [z_dir]
#purposes=['meh']
for purpose,rot_mat,offset_vec, propagation_vec \
in zip(purposes,rot_mats,offset_vecs,propagation_vecs):
# offset /= np.linalg.norm(offset) / 3e-6
# offset_vecs = [-offset]#[shear_offset]
# propagation_vecs = [z_dir]
# purposes=['meh']
for purpose, rot_mat, offset_vec, propagation_vec in zip(purposes, rot_mats, offset_vecs, propagation_vecs):
print('Testing', purpose)
Xs,Ys,Zs_0 = np.tensordot(rot_mat,Positions,axes=1)
Xs, Ys, Zs_0 = np.tensordot(rot_mat, Positions, axes=1)
new_basis = np.dot(rot_mat, basis)
Xs_prop, Ys_prop, Zs_prop = np.stack([Xs,Ys,Zs_0]) \
+ offset_vec[:,None,None]
Xs_prop, Ys_prop, Zs_prop = np.stack([Xs, Ys, Zs_0]) \
+ offset_vec[:, None, None]
print('Propagate Along', propagation_vec)
print('Propagate Along',propagation_vec)
if str(propagation_vec) == 'perp':
E0 = get_E(Xs,Ys,Zs_0, correct=False)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=False)
E0 = get_E(Xs, Ys, Zs_0, correct=False)
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=False)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,dtype=t.complex128)
E0.shape, new_basis, wavelength, offset_vec, dtype=t.complex128)
elif str(propagation_vec) == 'offset':
E0 = get_E(Xs,Ys,Zs_0, correct=True)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
E0 = get_E(Xs, Ys, Zs_0, correct=True)
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,
E0.shape, new_basis, wavelength, offset_vec,
dtype=t.complex128, propagate_along_offset=True)
else:
E0 = get_E(Xs,Ys,Zs_0, correct=True)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
E0 = get_E(Xs, Ys, Zs_0, correct=True)
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,
dtype=t.complex128, propagation_vector=propagation_vec)
E0.shape, new_basis, wavelength, offset_vec, dtype=t.complex128,
propagation_vector=propagation_vec)
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
if not check_equiv(Ez, Ez_t):
#if True:
# if True:
plt.close('all')
plt.figure()
plt.imshow(np.angle(E0))
@@ -443,70 +429,67 @@ def test_generalized_near_field():
plt.imshow(np.angle(Ez_t))
plt.title('Angle of numerically calculated Ez')
plt.figure()
plt.imshow(np.abs(Ez-Ez_t))#/np.max(np.abs(Ez)))
plt.imshow(np.abs(Ez - Ez_t))
plt.title('Magnitude of difference')
plt.show()
assert check_equiv(Ez, Ez_t)
Em0_t = propagators.inverse_near_field(t.as_tensor(Ez),asp).numpy()
assert check_equiv(E0,Em0_t)
Em0_t = propagators.inverse_near_field(t.as_tensor(Ez), asp).numpy()
assert check_equiv(E0, Em0_t)
print('Test Successful')
# One final test, to see if any of a few arbitrary rotations will
# change the predicted propagation if everything else is kept
# constant
Rrands = [stats.ortho_group.rvs(3) for i in range(3)]
Xs,Ys,Zs_0 = np.tensordot(Rboth,Positions,axes=1)
Xs, Ys, Zs_0 = np.tensordot(Rboth, Positions, axes=1)
new_basis = np.dot(rot_mat, basis)
offset_vec = shear_back_offset
propagation_vec = z_dir
E0 = get_E(Xs,Ys,Zs_0, correct=True)
E0 = get_E(Xs, Ys, Zs_0, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape, new_basis, wavelength,offset_vec,
dtype=t.complex128, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
E0.shape, new_basis, wavelength, offset_vec,
dtype=t.complex128, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
for Rrand in Rrands:
Xs,Ys,Zs_0 = np.tensordot(Rrand, np.tensordot(Rboth,Positions,axes=1),axes=1)
Xs, Ys, Zs_0 = np.tensordot(Rrand, np.tensordot(Rboth, Positions, axes=1), axes=1)
rot_offset = np.dot(Rrand, offset_vec)
rot_basis = np.dot(Rrand, new_basis)
rot_prop = np.dot(Rrand, propagation_vec)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape, rot_basis, wavelength, rot_offset,
dtype=t.complex128, propagation_vector=rot_prop)
Ez_rot_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
E0.shape, rot_basis, wavelength, rot_offset,
dtype=t.complex128, propagation_vector=rot_prop)
Ez_rot_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
assert np.max(np.abs(Ez_t - Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
assert np.max(np.abs(Ez_t-Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
def test_inverse_near_field():
x = (np.arange(800) - 400) * 1.5e-9
y = (np.arange(1200) - 600) * 1e-9
Ys,Xs = np.meshgrid(y,x)
Rs = np.sqrt(Xs**2+Ys**2)
wavelength = 3e-9 #nm
sigma = 20e-9 #nm
z = 1000e-9 #nm
Ys, Xs = np.meshgrid(y, x)
Rs = np.sqrt(Xs**2 + Ys**2)
w0 = np.sqrt(2)*sigma
wavelength = 3e-9 # nm
sigma = 20e-9 # nm
z = 1000e-9 # nm
w0 = np.sqrt(2) * sigma
E0 = np.exp(-Rs**2 / w0**2)
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
E0 = t.as_tensor(E0, dtype=t.complex128)
E_prop = propagators.near_field(E0, asp)
E0 = t.as_tensor(E0,dtype=t.complex128)
E_prop = propagators.near_field(E0,asp)
E_backprop = propagators.inverse_near_field(E_prop, asp)
# We just want to check that it actually is the inverse
assert t.all(t.isclose(E0,E_backprop))
assert t.all(t.isclose(E0, E_backprop))