diff --git a/CDTools/datasets/__init__.py b/CDTools/datasets/__init__.py index 30b0322..504647a 100644 --- a/CDTools/datasets/__init__.py +++ b/CDTools/datasets/__init__.py @@ -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: diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index d508d40..b66ef68 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -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) diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 1eb7802..8dc91a7 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -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 = [] diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index fe2246c..4f64030 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -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 diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 5e5deb1..937389c 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -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 diff --git a/CDTools/tools/analysis.py b/CDTools/tools/analysis.py index da24b7a..dc4ee22 100644 --- a/CDTools/tools/analysis.py +++ b/CDTools/tools/analysis.py @@ -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) diff --git a/CDTools/tools/data.py b/CDTools/tools/data.py index 8ee8849..6f955b6 100644 --- a/CDTools/tools/data.py +++ b/CDTools/tools/data.py @@ -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 diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py index 52326e8..961752d 100644 --- a/CDTools/tools/initializers.py +++ b/CDTools/tools/initializers.py @@ -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) diff --git a/CDTools/tools/losses.py b/CDTools/tools/losses.py index 2f92a45..7b92030 100644 --- a/CDTools/tools/losses.py +++ b/CDTools/tools/losses.py @@ -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) diff --git a/conda_requirements.txt b/conda_requirements.txt new file mode 100644 index 0000000..48fe15a --- /dev/null +++ b/conda_requirements.txt @@ -0,0 +1,8 @@ +numpy>=1.0 +scipy>=1.0 +matplotlib>=2.0 +python-dateutil +pytorch>=1.2.0 +h5py>=2.1 +pytest +sphinx diff --git a/docs/source/installation.rst b/docs/source/installation.rst index be1ab4f..7147748 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -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 `_ @@ -27,6 +39,7 @@ And has optional dependencies on * `pytest `_ * `sphinx `_ * `sphinx-argparse `_ + * `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. diff --git a/examples/example_reconstructions/gold_balls.pickle b/examples/example_reconstructions/gold_balls.pickle index 8db6ff1..2410051 100644 Binary files a/examples/example_reconstructions/gold_balls.pickle and b/examples/example_reconstructions/gold_balls.pickle differ diff --git a/examples/example_reconstructions/lab_ptycho.pickle b/examples/example_reconstructions/lab_ptycho.pickle index 395ae77..25e6ab5 100644 Binary files a/examples/example_reconstructions/lab_ptycho.pickle and b/examples/example_reconstructions/lab_ptycho.pickle differ diff --git a/examples/specular_ptycho.py b/examples/specular_ptycho.py index d85d7ba..f8615a3 100644 --- a/examples/specular_ptycho.py +++ b/examples/specular_ptycho.py @@ -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. diff --git a/setup.py b/setup.py index 4d70f23..f44fb81 100644 --- a/setup.py +++ b/setup.py @@ -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", diff --git a/tests/conftest.py b/tests/conftest.py index f1abff3..0894750 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 933bd70..c83e204 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -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. diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index 5045365..0732bbd 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -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 diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index d6f9b71..de80a6e 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -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