Fixed merging conflict - log scale in ptycho_2d_dataset.py

2026-09-09 13:02:41 +02:00 · 2019-10-03 15:25:27 -04:00
parent 61f8bf733e 69be77cc09
commit 8197eed50e
19 changed files with 144 additions and 49 deletions
@@ -32,8 +32,11 @@ import numpy as np
 import torch as t
 from copy import copy
 import h5py
-import pathlib
-
+try:
+    import pathlib
+except ImportError:
+    import pathlib2 as pathlib
+    
 from CDTools.tools import data as cdtdata
 from CDTools.tools import plotting
 from torch.utils import data as torchdata
@@ -98,7 +101,10 @@ class CDataset(torchdata.Dataset):
        self.wavelength = wavelength
        self.detector_geometry = copy(detector_geometry)
        if mask is not None:
-            self.mask = t.tensor(mask)
+            if isinstance(mask, t.Tensor):
+                self.mask = mask.detach().to(dtype=t.bool)
+            else:
+                self.mask = t.BoolTensor(mask)
        else:
            self.mask = None
        if background is not None:
@@ -3,7 +3,10 @@ import numpy as np
 import torch as t
 from copy import copy
 import h5py
-import pathlib
+try:
+    import pathlib
+except ImportError:
+    import pathlib2 as pathlib

 from CDTools.datasets import CDataset
 from CDTools.tools import data as cdtdata
@@ -62,7 +65,7 @@ class Ptycho2DDataset(CDataset):
        self.translations = t.tensor(translations)
        self.patterns = t.tensor(patterns)
        if self.mask is None:
-            self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.uint8)
+            self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.bool)
        self.mask.masked_fill_(t.isnan(t.sum(self.patterns,dim=(0,))),0)
        self.patterns.masked_fill_(t.isnan(self.patterns),0)

@@ -181,7 +184,7 @@ class Ptycho2DDataset(CDataset):
        cdtdata.add_ptycho_translations(cxi_file, self.translations)


-    def inspect(self):
+    def inspect(self, logarithmic=True):
        """Launches an interactive plot for perusing the data

        This launches an interactive plotting tool in matplotlib that
@@ -252,7 +255,10 @@ class Ptycho2DDataset(CDataset):
                cb1.ax.set_title('Integrated Intensity', size="medium", pad=5)
                cb1.ax.tick_params(labelrotation=20)

-                meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask)
+                if logarithmic:
+                    meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask)
+                else:
+                    meas = axes[1].imshow(meas_data * mask)

                cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.17,fraction=0.1)
                cb2.ax.tick_params(labelrotation=20)
@@ -280,7 +286,11 @@ class Ptycho2DDataset(CDataset):


                meas = axes[1].images[-1]
-                meas.set_data(np.log(meas_data) / np.log(10) * mask)
+                if logarithmic:
+                    meas.set_data(np.log(meas_data) / np.log(10) * mask)
+                else:
+                    meas.set_data(meas_data * mask)
+
                update_colorbar(meas)


@@ -131,13 +131,17 @@ class CDIModel(t.nn.Module):
        scheduler : torch.optim.lr_scheduler._LRScheduler
            Optional, a learning rate scheduler to use
        """
-
+        # First, calculate the normalization
+        normalization = 0
+        for inputs, patterns in data_loader:
+            normalization += t.sum(patterns).cpu().numpy()
+            
+            
        for it in range(iterations):
            loss = 0
            N = 0
            for inputs, patterns in data_loader:
                N += 1
-
                def closure():
                    optimizer.zero_grad()
                    sim_patterns = self.forward(*inputs)
@@ -151,14 +155,14 @@ class CDIModel(t.nn.Module):

                loss += optimizer.step(closure).detach().cpu().numpy()

-            loss /= N
+            loss /= normalization
            if scheduler is not None:
                scheduler.step(loss)
-
+                
            yield loss


-    def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False):
+    def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False, amsgrad=False):
        """Runs a round of reconstruction using the Adam optimizer
        
        This is generally accepted to be the most robust algorithm for use
@@ -184,7 +188,7 @@ class CDIModel(t.nn.Module):
                                           shuffle=True)

        # Define the optimizer
-        optimizer = t.optim.Adam(self.parameters(), lr = lr)
+        optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)


        # Define the scheduler
@@ -234,6 +238,46 @@ class CDIModel(t.nn.Module):
        return self.AD_optimize(iterations, data_loader, optimizer)


+    def SGD_optimize(self, iterations, dataset, batch_size=None,
+                     lr=0.01, momentum=0, dampening=0, weight_decay=0,
+                     nesterov=False):
+        """Runs a round of reconstruction using the SGDoptimizer
+        
+        This algorithm is often less stable that Adam, but it is simpler
+        and is the basic workhorse of gradience descent.
+
+        Parameters
+        ----------
+        iterations : int
+            How many epochs of the algorithm to run
+        dataset : CDataset
+            The dataset to reconstruct against
+        batch_size : int
+            Optional, the size of the minibatches to use
+        lr : float
+            Optional, the learning rate to use
+        momentum : float
+            Optional, the length of the history to use.
+        """
+        
+        # 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)
+
+
    # By default, the plot_list is empty
    plot_list = []
    
@@ -46,7 +46,7 @@ class FancyPtycho(CDIModel):
        if mask is None:
            self.mask = mask
        else:
-            self.mask = t.ByteTensor(mask)
+            self.mask = t.BoolTensor(mask)
        
        # We rescale the probe here so it learns at the same rate as the
        # object
@@ -175,7 +175,7 @@ class FancyPtycho(CDIModel):
        weights = t.ones(len(dataset))
        
        if hasattr(dataset, 'mask') and dataset.mask is not None:
-            mask = dataset.mask.to(t.uint8)
+            mask = dataset.mask.to(t.bool)
        else:
            mask = None

@@ -98,7 +98,7 @@ class SimplePtycho(CDIModel):


        if hasattr(dataset, 'mask') and dataset.mask is not None:
-            mask = dataset.mask.to(t.uint8)
+            mask = dataset.mask.to(t.bool)
        else:
            mask = None

@@ -148,15 +148,13 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):

    
    if correct_ramp:
-        # Need to check if this is actually working and, if noy, why not
-        center_freq = ip.centroid_sq(cmath.fftshift(t.fft(probe[0],2)),comp=True)
+        # Need to check if this is actually working and, if not, why not
+        center_freq = ip.centroid(cmath.cabssq(cmath.fftshift(t.fft(probe[0],2))))
        center_freq -= (t.tensor(probe[0].shape[:-1]) // 2).to(t.float32)
        center_freq /= t.tensor(probe[0].shape[:-1]).to(t.float32)
-
-
    
        Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].shape[1]]
-        probe_phase_ramp = cmath.expi(2*np.pi *
+        probe_phase_ramp = cmath.expi(2 * np.pi *
                                      (center_freq[0] * t.tensor(Is).to(t.float32) +
                                       center_freq[1] * t.tensor(Js).to(t.float32)))
        probe = cmath.cmult(probe, cmath.cconj(probe_phase_ramp))
@@ -165,7 +163,6 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
                                    (center_freq[0] * t.tensor(Is).to(t.float32) +
                                     center_freq[1] * t.tensor(Js).to(t.float32)))
        obj = cmath.cmult(obj, obj_phase_ramp)
-
    
    # Then, we set them to consistent absolute phases
    
@@ -475,7 +472,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
                          (im1.shape[1]//8)*3:(im1.shape[1]//8)*5]

    if nbins is None:
-        nbins = np.max(synth_obj[im_slice].shape) // 4
+        nbins = np.max(im1[im_slice].shape) // 4

    
    cor_fft = cmath.cmult(cmath.fftshift(t.fft(im1[im_slice],2)),
@@ -490,7 +487,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
    
    i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di))
    j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj))
-    
+
    Js,Is = np.meshgrid(j_freqs,i_freqs)
    Rs = np.sqrt(Is**2+Js**2)
    
@@ -278,7 +278,7 @@ def get_mask(cxi_file):
        mask = np.array(i1['detector_1/mask']).astype(np.uint32)
        mask_on = np.equal(mask,np.uint32(0))
        mask_has_signal = np.equal(mask,np.uint32(0x00001000))
-        return np.logical_or(mask_on,mask_has_signal).astype(np.uint8)
+        return np.logical_or(mask_on,mask_has_signal).astype(np.bool)
    else:
        return None

@@ -73,6 +73,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,

    # In some edge cases this shape can be smaller than the detector shape
    full_shape = t.max(full_shape, det_shape)
+
    
    if opt_for_fft:
        full_shape = t.Tensor([next_fast_len(dim) for dim in full_shape]).to(t.int32)
@@ -21,7 +21,11 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
    This function calculates the mean squared error between their
    associated amplitudes. Because this is not well defined for negative
    numbers, make sure that all the intensities are >0 before using this
-    loss.
+    loss. Note that this is actually a sum-squared error, because this
+    formulation makes it vastly simpler to compare error calculations
+    between reconstructions with different minibatch size. I hope to
+    find a better way to do this that is more honest with this
+    cost function, though.

    It can accept intensity and simulated intensity tensors of any shape
    as long as their shapes match, and the provided mask array can be
@@ -50,11 +54,11 @@ def amplitude_mse(intensities, sim_intensities, mask=None):

    if mask is None:
        return t.sum((t.sqrt(sim_intensities) -
-                      t.sqrt(intensities))**2) / intensities.view(-1).shape[0]
+                      t.sqrt(intensities))**2)
    else:
        masked_intensities = intensities.masked_select(mask)
        return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
-                      t.sqrt(masked_intensities))**2) / masked_intensities.shape[0]
+                      t.sqrt(masked_intensities))**2)


    
@@ -0,0 +1,8 @@
+numpy>=1.0
+scipy>=1.0
+matplotlib>=2.0
+python-dateutil
+pytorch>=1.2.0
+h5py>=2.1
+pytest
+sphinx
@@ -13,6 +13,18 @@ It is recommended that you clone the repository, rather than just downloading th
 Step 2: Install Dependencies
 ----------------------------

+The dependencies for CDTools can be installed, if you are managing your environment with anaconda, by running
+
+.. code:: bash
+	  
+   $ conda install --file conda_requirements.txt
+
+There are two optional dependencies which are not installed via this procedure - the dependency sphinx-argparse for building the docs, and the pathlib2 module that provides python 2 compatibility. These can either be installed manually via conda-forge, or otherwise they will be installed automatically by pip during the final installation step if needed.
+
+If you manage your environment with pip, all required packges should be installed automatically. The only thing to be aware of is that pytorch must be compiled with MKL support, and CUDA support if you would like to use the GPU. For this reason, using anaconda python is strongly recommended.
+
+For convenience, the full set of dependencies are noted below:
+   
 CDTools depends on the following packages:

   * `numpy <http://www.numpy.org>`_
@@ -27,6 +39,7 @@ And has optional dependencies on
   * `pytest <https://docs.pytest.org/>`_
   * `sphinx <https://www.sphinx-doc.org/>`_
   * `sphinx-argparse <https://sphinx-argparse.readthedocs.io>`_
+   * `pathlib2 <https://pypi.org/project/pathlib2/>`_
     
 All of these can be installed via pip or conda. Finally, CDTools is written to be python 2.7+ compatible, but is only actively tested on python 3.

@@ -42,7 +55,9 @@ To install in CDTools in developer mode (recommended, to allow any updates to be

 .. code:: bash
 	  
-   $ pip install -e .
+   $ pip install -e .[tests,docs]
+
+If you don't need to run the tests, or don't need to build the docs, you can omit the relevant option or options.

 If you prefer to use a tool other than pip, CDTools can be installed via any other package management tool that works with a setup.py file.

@@ -2,6 +2,7 @@ from __future__ import division, print_function, absolute_import

 import CDTools
 from matplotlib import pyplot as plt
+import numpy as np


 # This file is too large to be distributed via Github.
@@ -13,12 +13,18 @@ setuptools.setup(
    long_description_content_type="text/markdown",
    url="https://github.mit.edu/scattering/CDTools.git",
    install_requires=[
-        "numpy",
-        "scipy",
-        "matplotlib",
+        "numpy>=1.0",
+        "scipy>=1.0",
+        "matplotlib>=2.0",
        "python-dateutil",
-        "torch",
-        "h5py"],
+        "torch>=1.2.0", #1.2.0 introduced boolean tensors in a breaking way, we use the boolean tensors here for masking
+        "h5py>=2.1",
+        "pathlib2 ; python_version<'3.4'"],
+    extras_require={
+        'tests': ["pytest"],
+        'docs': ["sphinx","sphinx-argparse"],
+        ":python_version<'3.4'": ["pathlib2"],
+    },
    packages=setuptools.find_packages(),
    classifiers=[
        "Programming Language :: Python :: 3",
@@ -103,8 +103,8 @@ def ptycho_cxi_1():

    # Remember the format for the CXI file differs from the format used
    # internally
-    mask = np.zeros((256,256)).astype(np.uint32)
-    expected['mask'] = np.ones((256,256)).astype(np.uint8)
+    mask = np.zeros((256,256)).astype(np.int32)
+    expected['mask'] = np.ones((256,256)).astype(np.bool)
    d1f.create_dataset('mask',data=mask)

    # Create an initial background
@@ -272,7 +272,7 @@ def ptycho_cxi_3():
    # 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(np.uint8)
+    expected['mask'] = np.ones((256,256)).astype(np.bool)
    d1f.create_dataset('mask',data=mask)
    expected['dark'] = None
    
@@ -27,8 +27,8 @@ def test_CDataset_init():
    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)))
+
+    assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(np.bool))))
    assert dataset.entry_info == entry_info
    assert dataset.sample_info == sample_info
    assert dataset.wavelength == wavelength
@@ -110,7 +110,7 @@ def test_CDataset_to(ptycho_cxi_1):
    dataset = CDataset.from_cxi(ptycho_cxi_1[0])

    dataset.to(dtype=t.float32)
-    assert dataset.mask.dtype == t.uint8
+    assert dataset.mask.dtype == t.bool
    # If cuda is available, check that moving the mask to CUDA works.
    if t.cuda.is_available():
        dataset.to(device='cuda:0')
@@ -147,7 +147,7 @@ def test_Ptycho2DDataset_init():
                                detector_geometry=detector_geometry,
                                mask=mask)

-    assert t.all(t.eq(dataset.mask,t.tensor(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
@@ -241,7 +241,7 @@ 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.uint8
+    assert dataset.mask.dtype == t.bool
    assert dataset.patterns.dtype == t.float64
    assert dataset.translations.dtype == t.float64
    # If cuda is available, check that moving the mask to CUDA works.
@@ -8,7 +8,10 @@ import pytest
 import os
 import datetime
 import numbers
-from pathlib import Path
+try:
+    import pathlib
+except ImportError:
+    import pathlib2 as pathlib



@@ -16,17 +16,17 @@ def test_amplitude_mse():
    # And add some noise to it
    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(np.uint8)
+    mask = (np.random.rand(100,100) > 0.1).astype(np.bool)

    # First, test without a mask
    np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2)
-    np_result /=  data.size
+    #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))
+    #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))
@@ -38,7 +38,7 @@ def test_intensity_mse():
    # And add some noise to it
    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(np.uint8)
+    mask = (np.random.rand(100,100) > 0.1).astype(np.bool)


    # First, test without a mask
@@ -61,7 +61,7 @@ def test_poisson_ml():
    # And add some noise to it
    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(np.uint8)
+    mask = (np.random.rand(100,100) > 0.1).astype(np.bool)


    # First, test without a mask