diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..9214cb9 --- /dev/null +++ b/.flake8 @@ -0,0 +1,2 @@ +[flake8] +ignore = E501, W503 \ No newline at end of file diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 5cc6d9b..a24af2b 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -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 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..9beed2b --- /dev/null +++ b/.github/workflows/publish.yml @@ -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 \ No newline at end of file diff --git a/README.md b/README.md index 696ab4f..742b472 100644 --- a/README.md +++ b/README.md @@ -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). diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 517f1b6..a3ef41e 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -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..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 ---------------- diff --git a/docs/source/index.rst b/docs/source/index.rst index c6fcb8d..b8cd288 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -9,6 +9,7 @@ general datasets models + reconstructors tools/index indices_tables diff --git a/docs/source/installation.rst b/docs/source/installation.rst index d998f93..11c3aad 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -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 `_. +.. _`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 `_ >= 1.0 * `scipy `_ >= 1.0 * `matplotlib `_ >= 2.0 - * `pytorch `_ >= 1.9.0 + * `pytorch `_ >= 2.3.0 * `python-dateutil `_ * `h5py `_ >= 2.1 -Optional dependencies: +Optional dependencies for running tests: * `pytest `_ * `pooch `_ + +Optional dependencies for building docs: + * `sphinx `_ >= 4.3.0 * `sphinx-argparse `_ * `sphinx_rtd_theme `_ >= 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 `_ 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 `_ and we'll get to the bottom of it. diff --git a/docs/source/intro.rst b/docs/source/intro.rst index 3a2e8d4..22bc2c1 100644 --- a/docs/source/intro.rst +++ b/docs/source/intro.rst @@ -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 diff --git a/docs/source/reconstructors.rst b/docs/source/reconstructors.rst new file mode 100644 index 0000000..2e78a74 --- /dev/null +++ b/docs/source/reconstructors.rst @@ -0,0 +1,5 @@ +Reconstructors +============== + +.. automodule:: cdtools.reconstructors + :members: diff --git a/example_environment.yml b/example_environment.yml deleted file mode 100644 index 46767a2..0000000 --- a/example_environment.yml +++ /dev/null @@ -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 - diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py index 1cf5855..ebee68f 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -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) diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 4747666..4971975 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -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) diff --git a/examples/gold_ball_split.py b/examples/gold_ball_split.py index 8624b9f..9fc5b08 100644 --- a/examples/gold_ball_split.py +++ b/examples/gold_ball_split.py @@ -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()) diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 726af94..217d96b 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -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()) diff --git a/pyproject.toml b/pyproject.toml index b7fd46a..9ba1438 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 1c584c8..0000000 --- a/requirements.txt +++ /dev/null @@ -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 diff --git a/setup.py b/setup.py deleted file mode 100644 index 59d4600..0000000 --- a/setup.py +++ /dev/null @@ -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", - ], -) - diff --git a/src/cdtools/__init__.py b/src/cdtools/__init__.py index 9684207..19dc4c6 100644 --- a/src/cdtools/__init__.py +++ b/src/cdtools/__init__.py @@ -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 diff --git a/src/cdtools/_version.py b/src/cdtools/_version.py new file mode 100644 index 0000000..9163035 --- /dev/null +++ b/src/cdtools/_version.py @@ -0,0 +1 @@ +__version__ = "0.3.1.dev" diff --git a/src/cdtools/datasets/base.py b/src/cdtools/datasets/base.py index 9dc4801..3f8ec8c 100644 --- a/src/cdtools/datasets/base.py +++ b/src/cdtools/datasets/base.py @@ -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) diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py index 8b6f6d8..3825d6d 100644 --- a/src/cdtools/datasets/ptycho_2d_dataset.py +++ b/src/cdtools/datasets/ptycho_2d_dataset.py @@ -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 x 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, diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index b31061d..df347b6 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -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() diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index dba678a..df26ff3 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -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)) ] diff --git a/src/cdtools/reconstructors/__init__.py b/src/cdtools/reconstructors/__init__.py new file mode 100644 index 0000000..84a1a81 --- /dev/null +++ b/src/cdtools/reconstructors/__init__.py @@ -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 diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py new file mode 100644 index 0000000..e298bdb --- /dev/null +++ b/src/cdtools/reconstructors/adam.py @@ -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, + ) diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py new file mode 100644 index 0000000..1666814 --- /dev/null +++ b/src/cdtools/reconstructors/base.py @@ -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 diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py new file mode 100644 index 0000000..8907782 --- /dev/null +++ b/src/cdtools/reconstructors/lbfgs.py @@ -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) diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py new file mode 100644 index 0000000..cb1e26b --- /dev/null +++ b/src/cdtools/reconstructors/sgd.py @@ -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, + ) diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py index 0fd913e..9844edb 100644 --- a/src/cdtools/tools/analysis/analysis.py +++ b/src/cdtools/tools/analysis/analysis.py @@ -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 diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py index b69d0c0..89879dd 100644 --- a/src/cdtools/tools/data/data.py +++ b/src/cdtools/tools/data/data.py @@ -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 diff --git a/src/cdtools/tools/measurements/measurements.py b/src/cdtools/tools/measurements/measurements.py index e726149..c926568 100644 --- a/src/cdtools/tools/measurements/measurements.py +++ b/src/cdtools/tools/measurements/measurements.py @@ -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: diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py index abd8f84..0f570ea 100644 --- a/src/cdtools/tools/plotting/plotting.py +++ b/src/cdtools/tools/plotting/plotting.py @@ -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 + ')') diff --git a/tests/conftest.py b/tests/conftest.py index 51a58dc..f0faea5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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' diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index c9eaa1b..02893f9 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -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 diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py index f74dd0d..b6b1868 100644 --- a/tests/models/test_simple_ptycho.py +++ b/tests/models/test_simple_ptycho.py @@ -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 diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 6897991..4d17f0b 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -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, :]) diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py new file mode 100644 index 0000000..132a8c2 --- /dev/null +++ b/tests/test_reconstructors.py @@ -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 diff --git a/tests/test_version.py b/tests/test_version.py new file mode 100644 index 0000000..5b6a778 --- /dev/null +++ b/tests/test_version.py @@ -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" diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index 27c6a65..ad06757 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -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]) diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index e637509..7f2571c 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -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]) diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index 83c34d8..a6385cd 100644 --- a/tests/tools/test_image_processing.py +++ b/tests/tools/test_image_processing.py @@ -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) diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index 53c284a..5430271 100644 --- a/tests/tools/test_initializers.py +++ b/tests/tools/test_initializers.py @@ -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 diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index ae3d6bd..563d8f6 100644 --- a/tests/tools/test_interactions.py +++ b/tests/tools/test_interactions.py @@ -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) - diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index 633027d..b36b85d 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -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)) diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py index bfa06de..93ac72c 100644 --- a/tests/tools/test_measurements.py +++ b/tests/tools/test_measurements.py @@ -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, diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py index b377ef6..baca33d 100644 --- a/tests/tools/test_plotting.py +++ b/tests/tools/test_plotting.py @@ -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() diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index ec55140..d571b41 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -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))