diff --git a/.gitignore b/.gitignore index e21d376..2510940 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,6 @@ __pycache__/ /tutorial notebooks/experiment.dump /docs/html +/plots +/xrdfit/test.py +/xrdfit.egg-info diff --git a/README.md b/README.md index f0de3b5..df705f3 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,6 @@ To install as a Python module, type from the root directory. For developers, you should install in linked .egg mode using -`python -m pip install . develop` +`python setup.py develop` If you are using a Python virtual environment, you should activate this first before using the above commands. \ No newline at end of file diff --git a/setup.py b/setup.py index 012f4d9..23299c5 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,8 @@ 'dill', 'tqdm', 'scipy', - 'lmfit' + 'lmfit', + 'jupyter', ], extras_require={"documentation_compilation": "sphinx"} ) \ No newline at end of file diff --git a/xrdfit/plotting.py b/xrdfit/plotting.py index 4af64f8..caf1e9c 100644 --- a/xrdfit/plotting.py +++ b/xrdfit/plotting.py @@ -1,3 +1,5 @@ +import os +import pathlib from typing import Tuple, List, Union, TYPE_CHECKING import lmfit @@ -83,7 +85,7 @@ def plot_peak_params(peak_params: List["PeakParams"], x_range: Tuple[float, floa def plot_peak_fit(data: np.ndarray, cake_numbers: List[int], fit_result: lmfit.model.ModelResult, - fit_name: str): + fit_name: str, timestep: str = None, file_name: str = None): """Plot the result of a peak fit as well as the raw data.""" plt.figure(figsize=(8, 6)) @@ -100,9 +102,18 @@ def plot_peak_fit(data: np.ndarray, cake_numbers: List[int], fit_result: lmfit.m plt.xlabel(r'Two Theta ($^\circ$)') plt.ylabel('Intensity') plt.legend() + if timestep: + fit_name = f'Peak "{fit_name}" at t = {timestep}' plt.title(fit_name) plt.tight_layout() - plt.show() + if file_name: + file_name = pathlib.Path(file_name) + if not file_name.parent.exists(): + os.makedirs(file_name.parent) + plt.savefig(file_name) + else: + plt.show() + plt.close() def plot_parameter(data: np.ndarray, fit_parameter: str, peak_name: str, show_points: bool): diff --git a/xrdfit/spectrum_fitting.py b/xrdfit/spectrum_fitting.py index 13b64b9..3736ebc 100644 --- a/xrdfit/spectrum_fitting.py +++ b/xrdfit/spectrum_fitting.py @@ -61,6 +61,32 @@ def __str__(self) -> str: return f"PeakParams('{self.name}', {self.peak_bounds})" return f"PeakParams('{self.name}', {self.peak_bounds}, {self.maxima_bounds})" + def adjust_peak_bounds(self, fit_params: lmfit.Parameters): + """Adjust peak bounds to re-center the peak in the peak bounds.""" + centers = [] + for name in fit_params: + if "center" in name: + centers.append(fit_params[name].value) + center = sum(centers) / len(centers) + + bound_width = self.peak_bounds[1] - self.peak_bounds[0] + self.peak_bounds = (center - (bound_width / 2), center + (bound_width / 2)) + + def adjust_maxima_bounds(self, fit_params: lmfit.Parameters): + """Adjust maxima bounds to re-center the maximum in the maximum bounds.""" + peak_centers = [] + for param in fit_params: + if "center" in param: + peak_centers.append(fit_params[param].value) + + new_maxima_bounds = [] + for center, maximum_bounds in zip(peak_centers, self.maxima_bounds): + maximum_bound_width = maximum_bounds[1] - maximum_bounds[0] + lower_bound = center - maximum_bound_width / 2 + upper_bound = center + maximum_bound_width / 2 + new_maxima_bounds.append((lower_bound, upper_bound)) + self.maxima_bounds = new_maxima_bounds + class PeakFit: """An object containing data on the fit to a peak. @@ -74,12 +100,13 @@ def __init__(self, name: str): self.result: Union[None, lmfit.model.ModelResult] = None self.cake_numbers: List[int] = [] - def plot(self): + def plot(self, timestep: str = None, file_name: str = None): """ Plot the raw spectral data and the fit.""" if self.raw_spectrum is None: print("Cannot plot fit peak as fitting has not been done yet.") else: - plotting.plot_peak_fit(self.raw_spectrum, self.cake_numbers, self.result, self.name) + plotting.plot_peak_fit(self.raw_spectrum, self.cake_numbers, self.result, self.name, + timestep, file_name) class FitSpectrum: @@ -129,11 +156,14 @@ def plot(self, cakes_to_plot: Union[int, List[int]], x_range: Tuple[float, float plotting.plot_spectrum(data, cakes_to_plot, merge_cakes, show_points, x_range) plt.show() - def plot_fit(self, fit_name: str): + def plot_fit(self, fit_name: str, timestep: str = None, file_name: str = None): """Plot the result of a fit. - :param fit_name: The name of the fit to plot.""" + :param fit_name: The name of the fit to plot. + :param timestep: If provided, the timestep of the fit which will be added to the title. + :param file_name: If provided, the stub of the file name to write the plot to, if not + provided, the plot will be displayed on screen.""" fit = self.get_fit(fit_name) - fit.plot() + fit.plot(timestep, file_name) def plot_peak_params(self, peak_params: Union[PeakParams, List[PeakParams]], cakes_to_plot: Union[int, List[int]], @@ -319,7 +349,8 @@ def run_analysis(self, reuse_fits=False): for peak_fit, peak_params in zip(spectral_data.fitted_peaks, self.peak_params): if reuse_fits: peak_params.set_previous_fit(peak_fit.result.params) - self.adjust_peak_bounds(spectral_data.fitted_peaks) + peak_params.adjust_maxima_bounds(peak_fit.result.params) + peak_params.adjust_peak_bounds(peak_fit.result.params) print("Analysis complete.") @@ -365,14 +396,16 @@ def plot_fit_parameter(self, peak_name: str, fit_parameter: str, show_points=Fal plotting.plot_parameter(data, fit_parameter, peak_name, show_points) def plot_fits(self, num_timesteps: int = 5, peak_names: Union[List[str], str] = None, - timesteps: List[int] = None): + timesteps: List[int] = None, file_name: str = None): """Plot the calculated fits to the data. :param num_timesteps: The number of timesteps to plot fits for. The function will plot this many timesteps, evenly spaced over the whole dataset. This value is ignored if `timesteps` is specified. :param peak_names: The name of the peak to fit. If not specified, will plot all fitted peaks. - :param timesteps: A list of timesteps to plot the fits for.""" + :param timesteps: A list of timesteps to plot the fits for. + :param file_name: If provided, outputs the plot to an image file with filename as the image + stub.""" if timesteps is None: timesteps = self._calculate_timesteps(num_timesteps) @@ -380,7 +413,11 @@ def plot_fits(self, num_timesteps: int = 5, peak_names: Union[List[str], str] = peak_names = self.peak_names() for timestep in timesteps: for name in peak_names: - self.timesteps[timestep].plot_fit(name) + if file_name: + output_name = f"../plots/{file_name}_{name}_{timestep :04d}.png" + else: + output_name = None + self.timesteps[timestep].plot_fit(name, str(timestep), output_name) def _calculate_timesteps(self, num_timesteps: int) -> List[int]: """Work out which timesteps to plot.""" @@ -397,18 +434,6 @@ def save(self, file_name: str): dill.dump(self, output_file) print("Data successfully saved to dump file.") - def adjust_peak_bounds(self, fitted_peaks: List[PeakFit]): - """Adjust peak bounds to re-center the peak in the peak bounds.""" - for peak_fit, peak_param in zip(fitted_peaks, self.peak_params): - centers = [] - for name in peak_fit.result.params: - if "center" in name: - centers.append(peak_fit.result.params[name].value) - center = sum(centers) / len(centers) - - bound_width = peak_param.peak_bounds[1] - peak_param.peak_bounds[0] - peak_param.peak_bounds = (center - (bound_width / 2), center + (bound_width / 2)) - def get_stacked_spectrum(spectrum: np.ndarray) -> np.ndarray: """Take an number of observations from N different cakes and stack them vertically into a 2 diff --git a/xrdfit/test.py b/xrdfit/test.py deleted file mode 100644 index 99a8d99..0000000 --- a/xrdfit/test.py +++ /dev/null @@ -1,18 +0,0 @@ -import xrdfit.spectrum_fitting as spectrum_fitting -from xrdfit.spectrum_fitting import PeakParams - - -frame_time = 10 -file_stub = "../example_data/adc_041_7Nb_NDload_700C_15mms_*" -first_cake_angle = 90 -cakes_to_fit = [36, 1, 2] -peak_params = PeakParams('1', (2.8, 2.9)) -merge_cakes = True - -experiment = spectrum_fitting.FittingExperiment(frame_time, file_stub, first_cake_angle, - cakes_to_fit, peak_params, merge_cakes) - -experiment.run_analysis() -for peak_name in experiment.peak_names(): - for parameter in experiment.fit_parameters(peak_name): - experiment.plot_fit_parameter(peak_name, parameter) \ No newline at end of file