diff --git a/CDTools/__init__.py b/CDTools/__init__.py index fbe00ae..095a51f 100644 --- a/CDTools/__init__.py +++ b/CDTools/__init__.py @@ -1,3 +1,4 @@ from __future__ import division, print_function, absolute_import from CDTools import tools +from CDTools import datasets diff --git a/CDTools/datasets.py b/CDTools/datasets.py new file mode 100644 index 0000000..5bff38c --- /dev/null +++ b/CDTools/datasets.py @@ -0,0 +1,173 @@ +from __future__ import division, print_function, absolute_import +import numpy as np +import torch as t +from copy import copy + +# The naming overlap here is definitely going to get confusing. +from CDTools import tools +from torch.utils import data as torchdata + +__all__ = ['CDataset', 'Ptycho_2D_Dataset'] + +# +# I think that here should live a variety of datasets that all +# subclass the basic dataset class. They should be able to contain +# all the information that would be available in a CDI-type experiment. +# +# It's important for each kind of dataset to have tools to: +# +# * Pass too and from the GPU (where to store the data) +# * Be able to load itself intelligently from a cxi file +# * Be able to save itself intelligently to a cxi file +# +# And in addition, it should be relatively simple to initialize from +# data held as python structures +# +# The datasets need to be able to work just like pytorch datasets, +# in fact, they should subclass them, where they return the relevant +# information for each diffraction pattern when sliced. +# + + +# +# One question, though, is how to deal with information that models need +# to know as well as the datasets? The basic issue is that it would be +# nice to be able to write down a forward model and have it populate +# a dataset object with the simulated data. But you also want to be able +# to write a forward model that pulls the relevant information from a +# dataset that has been read from a real data's cxi file. +# +# One workaround would be to just have the things defined twice, but to +# make sure that every model can load/initialize itself from a dataset. +# If it also knows how to define itself from an initialization function +# or a python dataset, then it can easily also have a generic function to +# simulate the action of the model and create a dataset from that simulation. +# +# This is unrelated, but it will then be important to be able to save and load +# models easily from a predefined format. To be honest, this could just +# literally be a snippet of python code that redefines the specific model +# using only the base CDTools packages. That way a saved model could be a +# .py file that just creates the model on the spot. Not sure I love this +# though +# +# Is the trio of data, model, and reconstruction the correct thing though? +# Perhaps the reconstruction doesn't need to exist. The model can have +# a few more things involved in it, and the reconstructions could be a +# single function, for example, that only uses what's in the model. +# +# The issue with the automatic differentiation stuff is that it's history +# dependent in a sense. So you generate an optimization object and you take +# a step with it. +# +# Can all reconstruction algorithms be formulated in terms of an explicit +# forward step, backward step, and gradient with respect to the forward +# step? No, for example the position annealing step needs to do it's own +# thing. Okay, well then it could be possible to write an automatic +# differentiation sequence that can generate n steps of a given automatic +# differentiation solved. So given a model, you could just generically +# create n steps of automatic differentiation. But then also write an +# explicit ePIE step, etc. It could be interesting if the reconstruction +# algorithms defined generators which would run one step and then yield +# the loss at the end of that step. This would let you easily save out +# the loss but also would make it simple to include live plotting +# +# +# Then instead of a plan, you would just write a python script to follow +# the steps you wanted to follow. +# +# + + +# I think there should be a base CDataSet that has entry_info, sample_info, detector_info, and mask attributes. It can know how to load this data from a cxi file and write it out to a cxi file. It also can know how to pass this info to and from the GPU (not all of it needs to be passed in that way). + +# Then, I will write a few common datasets as a demonstration. First is +# A 2D CDI dataset, second is a 2D Ptycho 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): + + # 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 + + + def to(self,*args,**kwargs): + if mask is not None: + self.mask.to(*args,**kwargs) + + + @classmethod + def from_cxi(cls, cxi_file): + entry_info = tools.get_entry_info(cxi_file) + sample_info = tools.get_sample_info(cxi_file) + wavelength = tools.get_wavelength(cxi_file) + distance, basis, corner = tools.get_detector_geometry(cxi_file) + detector_geometry = {'distance' : distance, + 'basis' : basis, + 'corner' : corner} + mask = tools.get_mask(cxi_file) + return cls(entry_info = entry_info, + sample_info = sample_info, + wavelength=wavelength, + detector_geometry=detector_geometry, + mask=mask) + + + def to_cxi(self, cxi_file): + if self.entry_info is not None: + tools.add_entry_info(cxi_file, self.entry_info) + if self.sample_info is not None: + tools.add_sample_info(cxi_file, self.sample_info) + if self.wavelength is not None: + tools.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 + tools.add_detector(cxi_file, + self.detector_info['wavelength'], + self.detector_info['basis'], + corner = corner) + if self.mask is not None: + tools.add_mask(cxi_file, mask) + + + +class Ptycho_2D_Dataset(CDataset): + + def __init__(self,translations, patterns, **kwargs): + + super(CDataset,self).__init__(kwargs) + self.translations = t.tensor(translations) + self.patterns = t.tensor(patterns) + + + def to(self, *args, **kwargs): + super(CDataset,self).to(*args,**kwargs) + self.translations.to(*args, **kwargs) + self.patterns.to(*args, **kwargs) + + + def __len__(self): + return self.patterns.shape[0] + + def __get__(self,index): + return index, self.translations[index], self.patterns[index] diff --git a/CDTools/scripts/view_cxi.py b/CDTools/scripts/view_cxi.py index c8c437c..5f7a9c4 100644 --- a/CDTools/scripts/view_cxi.py +++ b/CDTools/scripts/view_cxi.py @@ -1,7 +1,6 @@ from __future__ import division, print_function, absolute_import import h5py import numpy as np -import math import os from PyQt5 import QtWidgets from matplotlib import pyplot as plt diff --git a/CDTools/tools/__init__.py b/CDTools/tools/__init__.py index 47a8cd2..64995ab 100644 --- a/CDTools/tools/__init__.py +++ b/CDTools/tools/__init__.py @@ -2,3 +2,4 @@ from __future__ import division, print_function, absolute_import from CDTools.tools import cmath from CDTools.tools import losses +from CDTools.tools import data