refactor datasets, wrote more docs, added one more test dataset

CDTools/datasets/__init__.py

@@ -0,0 +1,133 @@
+from __future__ import division, print_function, absolute_import
+import numpy as np
+import torch as t
+from copy import copy
+import h5py
+import pathlib
+
+from CDTools.tools import data as cdtdata
+from CDTools.tools import plotting
+from torch.utils import data as torchdata
+from matplotlib import pyplot as plt
+from matplotlib.widgets import Slider
+from matplotlib import ticker
+
+
+#
+# This loads and stores all the kinds of metadata that are common to
+# All different kinds of diffraction experiments
+# Other datasets can subclass this and not worry about loading and
+# saving that metadata.
+#
+
+class CDataset(torchdata.Dataset):
+
+    def __init__(self, entry_info=None, sample_info=None,
+                 wavelength=None,
+                 detector_geometry=None, mask=None,
+                 background=None):
+
+        # Force pass-by-value-like behavior to stop strangeness
+        self.entry_info = copy(entry_info)
+        self.sample_info = copy(sample_info)
+        self.wavelength = wavelength
+        self.detector_geometry = copy(detector_geometry)
+        if mask is not None:
+            self.mask = t.tensor(mask)
+        else:
+            self.mask = None
+        if background is not None:
+            self.background = t.Tensor(background)
+        else:
+            self.background = None
+    
+        self.get_as(device='cpu')
+
+            
+    def to(self,*args,**kwargs):
+        # The mask should always stay a uint8, but it should switch devices
+        mask_kwargs = copy(kwargs)
+        try:
+            mask_kwargs.pop('dtype')
+        except KeyError as r:
+            pass
+        
+        if self.mask is not None:
+            self.mask = self.mask.to(*args,**mask_kwargs)          
+        if self.background is not None:
+            self.background = self.background.to(*args,**kwargs)          
+
+
+    def get_as(self, *args, **kwargs):
+        self.get_as_args = (args, kwargs)
+
+
+    def __getitem__(self, index):
+        # Deals with loading to appropriate device/dtype, if
+        # specified via a call to get_as
+        inputs, outputs = self._load(index)
+        if hasattr(self, 'get_as_args'):
+            outputs = outputs.to(*self.get_as_args[0],**self.get_as_args[1])
+            moved_inputs = []
+            for inp in inputs:
+                try:
+                    moved_inputs.append(inp.to(*self.get_as_args[0],**self.get_as_args[1]) )
+                except:
+                    moved_inputs.append(inp)
+        else:
+            moved_inputs = inputs
+        return moved_inputs, outputs
+
+    
+    def _load(self, index):
+        # Internal function to load data
+        raise NotImplementedError()
+        
+            
+    @classmethod
+    def from_cxi(cls, cxi_file):
+
+        # 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:
+                return cls.from_cxi(f)
+        
+        entry_info = cdtdata.get_entry_info(cxi_file)
+        sample_info = cdtdata.get_sample_info(cxi_file)
+        wavelength = cdtdata.get_wavelength(cxi_file)
+        distance, basis, corner = cdtdata.get_detector_geometry(cxi_file)
+        detector_geometry = {'distance' : distance,
+                             'basis'    : basis,
+                             'corner'   : corner}
+        mask = cdtdata.get_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)
+    
+    
+    def to_cxi(self, cxi_file):
+        if self.entry_info is not None:
+            cdtdata.add_entry_info(cxi_file, self.entry_info)
+        if self.sample_info is not None:
+            cdtdata.add_sample_info(cxi_file, self.sample_info)
+        if self.wavelength is not None:
+            cdtdata.add_source(cxi_file, self.wavelength)
+        if self.detector_geometry is not None:
+            if 'corner' in self.detector_geometry:
+                corner = self.detector_geometry['corner']
+            else:
+                corner = None
+            cdtdata.add_detector(cxi_file,
+                               self.detector_geometry['distance'],
+                               self.detector_geometry['basis'],
+                               corner = corner)
+        if self.mask is not None:
+            cdtdata.add_mask(cxi_file, self.mask)
+        if self.background is not None:
+            cdtdata.add_dark(cxi_file, self.background)
+        
+
+from CDTools.datasets.ptycho_2d_dataset import Ptycho2DDataset

CDTools/datasets/ptycho_2d_dataset.py

@@ -5,6 +5,7 @@ from copy import copy
 import h5py
 import pathlib
 
+from CDTools.datasets import CDataset
 from CDTools.tools import data as cdtdata
 from CDTools.tools import plotting
 from torch.utils import data as torchdata
@@ -12,126 +13,7 @@ from matplotlib import pyplot as plt
 from matplotlib.widgets import Slider
 from matplotlib import ticker
 
-__all__ = ['CDataset', 'Ptycho_2D_Dataset']
-
-
-#
-# This loads and stores all the kinds of metadata that are common to
-# All different kinds of diffraction experiments
-# Other datasets can subclass this and not worry about loading and
-# saving that metadata.
-#
-
-class CDataset(torchdata.Dataset):
-
-    def __init__(self, entry_info=None, sample_info=None,
-                 wavelength=None,
-                 detector_geometry=None, mask=None,
-                 background=None):
-
-        # Force pass-by-value-like behavior to stop strangeness
-        self.entry_info = copy(entry_info)
-        self.sample_info = copy(sample_info)
-        self.wavelength = wavelength
-        self.detector_geometry = copy(detector_geometry)
-        if mask is not None:
-            self.mask = t.tensor(mask)
-        else:
-            self.mask = None
-        if background is not None:
-            self.background = t.Tensor(background)
-        else:
-            self.background = None
-    
-        self.get_as(device='cpu')
-
-            
-    def to(self,*args,**kwargs):
-        # The mask should always stay a uint8, but it should switch devices
-        mask_kwargs = copy(kwargs)
-        try:
-            mask_kwargs.pop('dtype')
-        except KeyError as r:
-            pass
-        
-        if self.mask is not None:
-            self.mask = self.mask.to(*args,**mask_kwargs)          
-        if self.background is not None:
-            self.background = self.background.to(*args,**kwargs)          
-
-
-    def get_as(self, *args, **kwargs):
-        self.get_as_args = (args, kwargs)
-
-
-    def __getitem__(self, index):
-        # Deals with loading to appropriate device/dtype, if
-        # specified via a call to get_as
-        inputs, outputs = self._load(index)
-        if hasattr(self, 'get_as_args'):
-            outputs = outputs.to(*self.get_as_args[0],**self.get_as_args[1])
-            moved_inputs = []
-            for inp in inputs:
-                try:
-                    moved_inputs.append(inp.to(*self.get_as_args[0],**self.get_as_args[1]) )
-                except:
-                    moved_inputs.append(inp)
-        else:
-            moved_inputs = inputs
-        return moved_inputs, outputs
-
-    
-    def _load(self, index):
-        # Internal function to load data
-        raise NotImplementedError()
-        
-            
-    @classmethod
-    def from_cxi(cls, cxi_file):
-
-        # 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:
-                return cls.from_cxi(f)
-        
-        entry_info = cdtdata.get_entry_info(cxi_file)
-        sample_info = cdtdata.get_sample_info(cxi_file)
-        wavelength = cdtdata.get_wavelength(cxi_file)
-        distance, basis, corner = cdtdata.get_detector_geometry(cxi_file)
-        detector_geometry = {'distance' : distance,
-                             'basis'    : basis,
-                             'corner'   : corner}
-        mask = cdtdata.get_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)
-    
-    
-    def to_cxi(self, cxi_file):
-        if self.entry_info is not None:
-            cdtdata.add_entry_info(cxi_file, self.entry_info)
-        if self.sample_info is not None:
-            cdtdata.add_sample_info(cxi_file, self.sample_info)
-        if self.wavelength is not None:
-            cdtdata.add_source(cxi_file, self.wavelength)
-        if self.detector_geometry is not None:
-            if 'corner' in self.detector_geometry:
-                corner = self.detector_geometry['corner']
-            else:
-                corner = None
-            cdtdata.add_detector(cxi_file,
-                               self.detector_geometry['distance'],
-                               self.detector_geometry['basis'],
-                               corner = corner)
-        if self.mask is not None:
-            cdtdata.add_mask(cxi_file, self.mask)
-        if self.background is not None:
-            cdtdata.add_dark(cxi_file, self.background)
-        
-
+__all__ = ['Ptycho2DDataset']
 
 #
 # This is the standard dataset for a 2D ptychography experiment,
@@ -139,11 +21,11 @@ class CDataset(torchdata.Dataset):
 # programs (only tested against SHARP)
 #
 
-class Ptycho_2D_Dataset(CDataset):
+class Ptycho2DDataset(CDataset):
 
     def __init__(self, translations, patterns, axes=None, *args, **kwargs):
 
-        super(Ptycho_2D_Dataset,self).__init__(*args, **kwargs)
+        super(Ptycho2DDataset,self).__init__(*args, **kwargs)
         self.axes = copy(axes)
         self.translations = t.tensor(translations)
         self.patterns = t.tensor(patterns)
@@ -162,7 +44,7 @@ class Ptycho_2D_Dataset(CDataset):
 
 
     def to(self, *args, **kwargs):
-        super(Ptycho_2D_Dataset,self).to(*args,**kwargs)
+        super(Ptycho2DDataset,self).to(*args,**kwargs)
         self.translations = self.translations.to(*args, **kwargs)
         self.patterns = self.patterns.to(*args, **kwargs)
         
@@ -198,7 +80,7 @@ class Ptycho_2D_Dataset(CDataset):
     
 
     def to_cxi(self, cxi_file):
-        super(Ptycho_2D_Dataset,self).to_cxi(cxi_file)
+        super(Ptycho2DDataset,self).to_cxi(cxi_file)
         cdtdata.add_data(cxi_file, self.patterns, axes=self.axes)
         cdtdata.add_ptycho_translations(cxi_file, self.translations)
 

CDTools/models/fancy_ptycho.py

@@ -2,7 +2,7 @@ from __future__ import division, print_function, absolute_import
 
 import torch as t
 from CDTools.models import CDIModel
-from CDTools.datasets import Ptycho_2D_Dataset
+from CDTools.datasets import Ptycho2DDataset
 from CDTools import tools
 from CDTools.tools import cmath
 from CDTools.tools import plotting as p
@@ -316,7 +316,7 @@ class FancyPtycho(CDIModel):
         data = self.forward(indices, translations)
 
         # And finally, we make the dataset
-        return Ptycho_2D_Dataset(translations, data,
+        return Ptycho2DDataset(translations, data,
                                  entry_info = entry_info,
                                  sample_info = sample_info,
                                  wavelength=wavelength,

CDTools/models/pinhole_plane_ptycho.py

@@ -2,7 +2,7 @@ from __future__ import division, print_function, absolute_import
 
 import torch as t
 from CDTools.models import CDIModel
-from CDTools.datasets import Ptycho_2D_Dataset
+from CDTools.datasets import Ptycho2DDataset
 from CDTools import tools
 from CDTools.tools import cmath
 from CDTools.tools import plotting as p
@@ -331,7 +331,7 @@ class PinholePlanePtycho(CDIModel):
         data = self.forward(indices, translations)
 
         # And finally, we make the dataset
-        return Ptycho_2D_Dataset(translations, data,
+        return Ptycho2DDataset(translations, data,
                                  entry_info = entry_info,
                                  sample_info = sample_info,
                                  wavelength=wavelength,

CDTools/models/simple_ptycho.py

@@ -2,7 +2,7 @@ from __future__ import division, print_function, absolute_import
 
 import torch as t
 from CDTools.models import CDIModel
-from CDTools.datasets import Ptycho_2D_Dataset
+from CDTools.datasets import Ptycho2DDataset
 from CDTools import tools
 from CDTools.tools import plotting as p
 from copy import copy
@@ -177,7 +177,7 @@ class SimplePtycho(CDIModel):
         data = self.forward(indices, translations)
 
         # And finally, we make the dataset
-        return Ptycho_2D_Dataset(translations, data,
+        return Ptycho2DDataset(translations, data,
                                  entry_info = entry_info,
                                  sample_info = sample_info,
                                  wavelength=wavelength,

docs/source/datasets.rst

@@ -4,4 +4,6 @@ Datasets
 .. automodule:: CDTools.datasets
    :members:
 
+
+
    

docs/source/general.rst

@@ -1,11 +1,32 @@
 General Reference
 =================
 
+The full documentation contains the details of all the functions, classes, etc. contained in CDTools - however, there are a few issues which cut across the various functions and are important enough to note in one place. This includes some definitions which are used across the packages, as well as a few conventions which are important to understand before writing code with CDTools. Because of that, it is recommended to read through this general reference page before diving into the reference documentation.
 
-A note about the meaning of "array" as a type
 
-A note about SI units in the package
+Arrays, Numpy, and Pytorch
+--------------------------
 
-Notes about special things like the probe-convention translations and the transposed basis arrays
+By necessity, CDTools operates using a mixture of numpy arrays and pytorch tensors. This, unfortunately, sucks. Certain functions - such as the various tools designed to be used in ptychography models - only accept pytorch tensors. Other functions - mostly analysis functions - only work with numpy arrays. Where possible (for example, in many analysis functions), input is accepted in either format.
+
+Functions that can accept either pytorch tensors or numpy arrays will have the type for the relevant inputs documented as "array", rather than "np.ndarray" or "torch.Tensor". In general, these functions will either accept all "array" type inputs as numpy arrays, or all as torch tensors. They will then return a result in a format matching that of the inputs. While many of these functions will work with mixed numpy/pytorch input, the output behavior of these functions is not, in general, defined for such a case, so it is heavily discouraged.
+
+Finally, it is important to remember always that, since pytorch does not have complex number support, complex-valued tensors in pytorch are always represented by a tensor with a trailing dimension of length 2. All complex arithmetic operations are defined in :code:`tools.cmath`, but do not forget to use them!
+
+
+Unit Conventions
+----------------
+
+All physical units used everywhere are SI units, with no exception. Every unit of length is assumed to be meters, all units of energy are Joules, etc. This matches the .CXI file specification and provides easy interoperability.
+
+
+Deviations from CXI Conventions
+-------------------------------
+
+There are several intentional deviations between the conventions used for the ptychography dataset classes and the .CXI file spec.
+
+First, all probe translations are stored as translations of the probe over the object. We have found that this more closely maps onto the actual way that most ptychography experiments are run, and it is a more natural choice to use when cropping out a section of the object to multiply the probe with. However, .CXI files are defined to store translations of the object, not the probe - thus, all translations are inverted when read from a .CXI file, and again when written out to a .CXI file
+
+Second, the convention for storing bases in CDTools is transposed from the convention used in .CXI files. This means that the basis vectors are column vectors, which we find more natural when they are used to define a coordinate system. CDTools will accept a basis stored either in this (nonstandard) format in a .CXI file, or a basis stored in the appropriate format, and any files saved using the built-in tools will produce compliant .CXI files.
 
 

docs/source/index.rst

@@ -21,11 +21,11 @@ CDTools is a python library for ptychography and CDI reconstructions, using an A
 
    # imports
    from matplotlib import pyplot as plt
-   from CDTools.datasets import Ptycho_2D_Dataset
+   from CDTools.datasets import Ptycho2DDataset
    from CDTools.models import SimplePtycho
 		
    # Load the file
-   dataset = Ptycho_2D_Dataset.from_cxi('ptycho_data.cxi')
+   dataset = Ptycho2DDataset.from_cxi('ptycho_data.cxi')
 
    # Generate a model from the data
    model = SimplePtycho.from_dataset(dataset)
@@ -40,7 +40,7 @@ CDTools is a python library for ptychography and CDI reconstructions, using an A
    plt.show()
 
 
-CDTools makes it simple to load and inspect from data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it is straightforward to program new models for AD ptychography, which can then be used right away from the same scripting framework.
+CDTools makes it simple to load and inspect data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it includes a bunch of modular functions for AD ptychography, which can then be used right away from the same scripting framework.
 
 The high-level interface to CDTools is built on a lower level "three-legged stool". This consists of tools to access stored data, tools to visualize data and reconstructions, and tools that implement basic operations relevant to coherent diffraction. All of these tools can be used directly alongside the high-level interface, when needed.
 

docs/source/tutorial.rst

@@ -1,2 +1,114 @@
 Tutorial
 ========
+
+This tutorial builds on the examples, leading to a more complete understanding of how to use CDTools and - importantly - how to extend it. First, we will cover the details of writing a useful reconstruction script for a particular experiment. Next, we will discuss how to implement a new dataset type for different kinds of coherent diffraction. Finally, we will go over how to make new models to cover specific types of ptychography which aren't described by any of the built-in models.
+
+
+Reconstruction Scripts
+----------------------
+
+In this section, we will write a script to run a reconstruction on a dataset collected from our benchtop optical ptychography playground. This mirrors very closely the reconstruction examples, however I encourage everyone to follow along, writing this script out line-by-line, to help you learn more permanently the process of writing a custom reconstruction script.
+
+Our first step will be creating the file and filling out the boilerplate: All the imports we'll need.
+
+.. code-block:: python
+
+   from __future__ import division, print_function, absolute_import
+
+   import CDTools
+   from matplotlib import pyplot as plt
+   import pickle
+
+
+You can always import more libraries, like numpy, or pytorch, or pandas, or what have you, as needed. Next, we load the dataset and give it a look-over
+
+
+.. code-block:: python
+
+   filename = 'example_data/lab_ptycho_data.cxi'
+   dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
+
+   dataset.inspect()
+   plt.show()
+
+Now, run this script! You should see a window pop up, showing a nanomap of the integrated intensities at each scan point on one side, and an individual diffraction pattern on the other. You can then click around to make sure that everything is in order.
+
+Now that we know we have the data loaded and it looks good, we can go ahead and comment out the dataset inspecting code, and move on to creating a model. It's usually a good idea to start by loading a standard :code:`FancyPtycho` model without any special changes, sending it to the GPU.
+
+.. code-block:: python
+		
+   model = CDTools.models.FancyPtycho.from_dataset(dataset)
+   model.to(device='cuda')
+   dataset.get_as(device='cuda')
+
+
+We then try a basic Adam reconstruction with this model, with no changes to the defaults, to see how it works. 
+
+.. code-block:: python
+
+   for i, loss in enumerate(model.Adam_optimize(50, dataset)):
+       model.inspect(dataset)
+       print(i,loss)
+
+   model.compare()
+   plt.show()
+
+
+It is worth noting here exactly how this code is working. The reconstruction methods are actually returning generators. Generators in pythons are objects that work like lists, or tuples, but have to be read out one item at a time, from left to right. The catch is that, instead of just reading out objects from a list, they can run arbitrary code each time they are asked for the "next" item.
+
+In CDTools, every reconstruction method will return a generator. Whenever the generator is asked for the next item, it runs a single epoch of the reconstructionalgorithm, and then returns the average loss over that epoch as that next item. This allows the execution of the reconstruction algorithm to pause once every epoch, allowing some time for the user to run a small snippet of code to inspect how the reconstruction is coming along.
+
+From the end user perspective, all this means is: follow the format above, or more generally put the :code:`model.Adam_optimize(n, dataset)` call anywhere that you would feel comfortable putting a call to :code:`range(n)` - list comprehensions, for loops, etc.
+
+Once we run this, we can take a look at the result. What we see is pretty good, but we can see that there are some issues with the reconstruction near the edge, and the probe itself seems to be larger than the "stage" on which we're reconstructing it. So, we can make two tweaks to this code in response. First, we increase the oversampling ratio, which doubles the size of the stage (this often can cause other issues as well, but generally works well in situations like this where the probe is honestly too large.
+
+.. code-block:: python
+
+   model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=2)
+
+
+And secondly, we note that there don't seem to be any errors with the positioning. So we can just not reconstruct the probe positions, knowing that the initial guesses are already accurate enough. We can do this by writing the following line, just before we run the reconstruction for loop.
+
+.. code-block:: python
+		
+   model.translation_offsets.requires_grad = False
+
+What is going on here is that, when running the optimization algorithm, pytorch will automatically calculate gradients for and then optimize over a number of parameters defined in the model - this includes parameters like :code:`model.probe`, :code:`model.obj`, :code:`model.background`, etc. We can tell pytorch to stop calculating gradients for (and stop updating) any of these parameters by setting their :code:`requires_grad` property to :code:`False`.
+
+After running this reconstruction, we can see that we're getting a little improvement (and a larger field of view) by using oversampling, but out in the corners we're nucleating extra probes! We can fix this by adding a probe support - that is, declating that the probe has to be defined only within a certain box. This can be done most easily with an argument to the model constructor:
+
+.. code-block:: python
+   
+   model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=2,
+                                                   probe_support_radius=90)
+
+
+It also seems like we need a few more iterations to finish converging, so we up the iteration count to 100.
+
+.. code-block:: python
+
+   for i, loss in enumerate(model.Adam_optimize(100, dataset)):
+						   
+
+Now we expect to get a nice reconstruction, so we can save the data. You can save the data in any form you like, once the relevant information is extracted from the model and put into a dictionary. The standard method for saving out this information is as follows:
+
+.. code-block:: python
+
+   with open('example_reconstructions/lab_ptycho.pickle', 'wb') as f:
+       pickle.dump(model.save_results(dataset),f)
+
+This is usually placed before the call to :code:`plt.show()`, to make sure that if the user manually exits the program once all the plot windows are opened, the data will still have been saved.
+
+Now, your file should match the example file in examples/lab_ptycho_data.py. 
+
+
+Datasets
+--------
+
+In this section, we will write a dataset class that can be used to import a nonstandard kind of ptychographic data where the beam is scanned longitudinally though the sample rather than rastered across the sample laterally. At the end of this tutorial, we will have written the class defined in 
+
+
+Models
+------
+
+In this section, we will write a model to perform a reconstruction on the axial scanning ptychography d

examples/MIT_BNL_logo.py

@@ -7,7 +7,7 @@ from matplotlib import pyplot as plt
 # This file is too large to be distributed via Github.
 # Please contact Abe Levitan (alevitan@mit) if you would like access
 filename = '/media/Data Bank/CSX_6_17/Processed_CXIs/79511_p.cxi'
-dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 # In this dataset, the edges of the patterns are masked off anyway
 # We can easily just remove this data instead of leaving it to float.

examples/ensemble_reconstruction.py

@@ -5,7 +5,7 @@ import pickle
 
 # Load the data
 filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
-dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 results = []
 

examples/gold_ball_ptycho.py

@@ -6,7 +6,7 @@ import pickle
 
 # First, we load an example dataset from a .cxi file
 filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
-dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 # Next, we create a ptychography model from the dataset
 # Note that we explicitly as for two incoherent probe modes

examples/inspect_dataset.py

@@ -5,7 +5,7 @@ from matplotlib import pyplot as plt
 
 # First, we load an example dataset from a .cxi file
 filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
-dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 # And we take a look at the data
 dataset.inspect()

examples/lab_ptycho_data.py

@@ -0,0 +1,28 @@
+from __future__ import division, print_function, absolute_import
+
+import CDTools
+from matplotlib import pyplot as plt
+import pickle
+
+filename = 'example_data/lab_ptycho_data.cxi'
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
+
+# dataset.inspect()
+# plt.show()
+
+model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=2,
+                                                probe_support_radius=90)
+model.to(device='cuda')
+dataset.get_as(device='cuda')
+
+
+model.translation_offsets.requires_grad = False
+for i, loss in enumerate(model.Adam_optimize(50, dataset)):
+    model.inspect(dataset)
+    print(i,loss)
+    
+with open('example_reconstructions/lab_ptycho.pickle', 'wb') as f:
+    pickle.dump(model.save_results(dataset),f)
+
+model.compare(dataset)
+plt.show()

examples/simple_ptycho.py

@@ -5,7 +5,7 @@ from matplotlib import pyplot as plt
 
 # First, we load an example dataset from a .cxi file
 filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
-dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 # Next, we create a ptychography model from the dataset
 model = CDTools.models.SimplePtycho.from_dataset(dataset)

examples/specular_ptycho.py

@@ -7,7 +7,7 @@ from matplotlib import pyplot as plt
 # This file is too large to be distributed via Github.
 # Please contact Abe Levitan (alevitan@mit) if you would like access
 filename = '/media/Data Bank/CSX_10_18/Processed_CXIs/110531_p.cxi'
-dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 
 # This model definition includes lots of tweaks, described below.

tests/test_datasets.py

@@ -125,7 +125,7 @@ def test_CDataset_to(ptycho_cxi_1):
 #
 
 
-def test_Ptycho_2D_Dataset_init():
+def test_Ptycho2DDataset_init():
     entry_info = {'start_time': datetime.datetime.now(),
                   'title' : 'A simple test'}
     sample_info = {'name': 'A test sample',
@@ -140,7 +140,7 @@ def test_Ptycho_2D_Dataset_init():
     patterns = np.random.rand(20,256,256)
     translations = np.random.rand(20,3)
     
-    dataset = Ptycho_2D_Dataset(translations, patterns,
+    dataset = Ptycho2DDataset(translations, patterns,
                                 entry_info=entry_info,
                                 sample_info=sample_info,
                                 wavelength=wavelength,
@@ -157,9 +157,9 @@ def test_Ptycho_2D_Dataset_init():
 
 
 
-def test_Ptycho_2D_Dataset_from_cxi(test_ptycho_cxis):
+def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
-        dataset = Ptycho_2D_Dataset.from_cxi(cxi)
+        dataset = Ptycho2DDataset.from_cxi(cxi)
 
         # The entry metadata loaded
         for key in expected['entry metadata']:
@@ -195,17 +195,17 @@ def test_Ptycho_2D_Dataset_from_cxi(test_ptycho_cxis):
 
 
             
-def test_Ptycho_2D_Dataset_to_cxi(test_ptycho_cxis, tmp_path):
+def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
     for cxi, expected in test_ptycho_cxis:
         print('loading dataset')
-        dataset = Ptycho_2D_Dataset.from_cxi(cxi)
+        dataset = Ptycho2DDataset.from_cxi(cxi)
         print('dataset mask is type', dataset.mask.dtype)
-        with cdtdata.create_cxi(tmp_path / 'test_Ptycho_2D_Dataset_to_cxi.cxi') as f:
+        with cdtdata.create_cxi(tmp_path / 'test_Ptycho2DDataset_to_cxi.cxi') as f:
             dataset.to_cxi(f)
 
         # Now we have to check that all the stuff was written
-        with h5py.File(tmp_path / 'test_Ptycho_2D_Dataset_to_cxi.cxi', 'r') as f:
-            read_dataset = Ptycho_2D_Dataset.from_cxi(f)
+        with h5py.File(tmp_path / 'test_Ptycho2DDataset_to_cxi.cxi', 'r') as f:
+            read_dataset = Ptycho2DDataset.from_cxi(f)
 
         assert dataset.entry_info == read_dataset.entry_info
 
@@ -237,8 +237,8 @@ def test_Ptycho_2D_Dataset_to_cxi(test_ptycho_cxis, tmp_path):
         assert t.allclose(dataset.translations, read_dataset.translations)
 
     
-def test_Ptycho_2D_Dataset_to(ptycho_cxi_1):
-    dataset = Ptycho_2D_Dataset.from_cxi(ptycho_cxi_1[0])
+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
@@ -254,9 +254,9 @@ def test_Ptycho_2D_Dataset_to(ptycho_cxi_1):
 
 
         
-def test_Ptycho_2D_Dataset_ops(ptycho_cxi_1):
+def test_Ptycho2DDataset_ops(ptycho_cxi_1):
     cxi, expected = ptycho_cxi_1
-    dataset = Ptycho_2D_Dataset.from_cxi(cxi)
+    dataset = Ptycho2DDataset.from_cxi(cxi)
     dataset.get_as('cpu')
 
     assert len(dataset) == expected['data'].shape[0]
@@ -266,9 +266,9 @@ def test_Ptycho_2D_Dataset_ops(ptycho_cxi_1):
     assert t.allclose(pattern, t.tensor(expected['data'][3,:,:]))
 
 
-def test_Ptycho_2D_Dataset_get_as(ptycho_cxi_1):
+def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
     cxi, expected = ptycho_cxi_1
-    dataset = Ptycho_2D_Dataset.from_cxi(cxi)
+    dataset = Ptycho2DDataset.from_cxi(cxi)
     if t.cuda.is_available():
         dataset.get_as('cuda:0')
         assert len(dataset) == expected['data'].shape[0]

tests/tools/test_initializers.py

@@ -2,7 +2,7 @@ from __future__ import division, print_function, absolute_import
 
 from CDTools.tools import initializers
 from CDTools.tools import cmath
-from CDTools.datasets import Ptycho_2D_Dataset
+from CDTools.datasets import Ptycho2DDataset
 import numpy as np
 import torch as t
 
@@ -112,7 +112,7 @@ def test_gaussian():
 
 def test_gaussian_probe(ptycho_cxi_1):
     
-    dataset = Ptycho_2D_Dataset.from_cxi(ptycho_cxi_1[0])
+    dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
 
     det_basis = t.Tensor(dataset.detector_geometry['basis'])
     det_shape = t.Size(dataset.patterns.shape[-2:])
@@ -173,7 +173,7 @@ def test_SHARP_style_probe(ptycho_cxi_1):
     # This code will probably change and honestly it doesn't need to
     # be exactly the final thing. So just test that the function doesn't
     # throw an error.
-    dataset = Ptycho_2D_Dataset.from_cxi(ptycho_cxi_1[0])
+    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