From 2d2f0ec5e482bd629ecf45d91004f63f1f1c6be3 Mon Sep 17 00:00:00 2001 From: Artur Glavic Date: Mon, 6 Oct 2025 17:59:09 +0200 Subject: [PATCH] Add plot command line option and method for projections --- eos/{normalisation.py => normalization.py} | 1 - eos/options.py | 25 +++++++++++++++- eos/projection.py | 34 +++++++++++++++++++++- eos/reduction.py | 30 +++++++++++++++++-- tests/test_full_analysis.py | 2 +- 5 files changed, 86 insertions(+), 6 deletions(-) rename eos/{normalisation.py => normalization.py} (98%) diff --git a/eos/normalisation.py b/eos/normalization.py similarity index 98% rename from eos/normalisation.py rename to eos/normalization.py index b1e7b33..53cf570 100644 --- a/eos/normalisation.py +++ b/eos/normalization.py @@ -26,7 +26,6 @@ class LZNormalisation: norm_lz, _, _ = np.histogram2d(lamda_e, detZ_e, bins=(grid.lamda(), grid.z())) norm_lz = np.where(norm_lz>2, norm_lz, np.nan) if normalisationMethod==NormalisationMethod.direct_beam: - # TODO: move flipping to projection self.norm = np.flip(norm_lz, 1) else: # correct for reference sm reflectivity diff --git a/eos/options.py b/eos/options.py index c9c8d30..ec6f470 100644 --- a/eos/options.py +++ b/eos/options.py @@ -434,6 +434,13 @@ class OutputFomatOption(StrEnum): Rlt = "Rlt" +class PlotColormaps(StrEnum): + gist_ncar = "gist_ncar" + viridis = "viridis" + inferno = "inferno" + gist_rainbow = "gist_rainbow" + nipy_spectral = "nipy_spectral" + @dataclass class OutputConfig(ArgParsable): outputFormats: List[OutputFomatOption] = field( @@ -460,6 +467,22 @@ class OutputConfig(ArgParsable): 'help': '?', }, ) + plot: bool = field( + default=False, + metadata={ + 'group': 'output', + 'help': 'show matplotlib graphs of results', + }, + ) + + plot_colormap: PlotColormaps = field( + default=PlotColormaps.gist_ncar, + metadata={ + 'short': 'pcmap', + 'group': 'output', + 'help': 'matplotlib colormap used in lambda-theta graphs when plotting', + }, + ) def _output_format_list(self, outputFormat): format_list = [] @@ -531,7 +554,7 @@ class EOSConfig: otpt += f' -of {" ".join(self.output.outputFormats)}' mask = '' - # TODO: Check if you want these parameters for the case of default call + mask += f' -y {" ".join(str(ii) for ii in self.experiment.yRange)}' mask += f' -l {" ".join(str(ff) for ff in self.experiment.lambdaRange)}' mask += f' -t {" ".join(str(ff) for ff in self.reduction.thetaRange)}' diff --git a/eos/projection.py b/eos/projection.py index e6a7c5b..894208d 100644 --- a/eos/projection.py +++ b/eos/projection.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from .event_data_types import EventDatasetProtocol from .instrument import Detector, LZGrid -from .normalisation import LZNormalisation +from .normalization import LZNormalisation @dataclass class ProjectedReflectivity: @@ -68,6 +68,13 @@ class ProjectedReflectivity: self.R -= R self.dR = np.sqrt(self.dR**2+dR**2) + def plot(self, **kwargs): + from matplotlib import pyplot as plt + plt.errorbar(self.Q, self.R, xerr=self.dQ, yerr=self.dR, **kwargs) + plt.yscale('log') + plt.xlabel('Q / $\\AA^{-1}$') + plt.ylabel('R') + class LZProjection: grid: LZGrid lamda: np.ndarray @@ -283,3 +290,28 @@ class LZProjection: right_list = cls.devide_bin(lambda_e[right_region], position_e[right_region], lamda_edges[split_idx:], dimension) return left_list+right_list + + def plot(self, **kwargs): + from matplotlib import pyplot as plt + from matplotlib.colors import LogNorm + + if 'colorbar' in kwargs: + cmap=True + del(kwargs['colorbar']) + else: + cmap=False + + if self.is_normalized: + if not 'norm' in kwargs: + kwargs['norm'] = LogNorm(2e-3, 2.0) + plt.pcolormesh(self.lamda, self.alphaF, self.data.ref, **kwargs) + if cmap: + plt.colorbar(label='R') + else: + if not 'norm' in kwargs: + kwargs['norm'] = LogNorm() + plt.pcolormesh(self.lamda, self.alphaF, self.data.I, **kwargs) + if cmap: + plt.colorbar(label='I / cpm') + plt.xlabel('$\\lambda$ / $\\AA$') + plt.ylabel('$\\Theta$ / °') diff --git a/eos/reduction.py b/eos/reduction.py index 3a2a6ec..5c0700b 100644 --- a/eos/reduction.py +++ b/eos/reduction.py @@ -10,7 +10,7 @@ from .header import Header from .path_handling import PathResolver from .options import EOSConfig, IncidentAngle, MonitorType, NormalisationMethod from .instrument import Detector, LZGrid -from .normalisation import LZNormalisation +from .normalization import LZNormalisation from . import event_handling as eh, event_analysis as ea from .projection import LZProjection @@ -113,6 +113,13 @@ class AmorReduction: if 'Rlt.ort' in self.output_config.outputFormats: self.save_Rtl() + if self.output_config.plot: + import matplotlib.pyplot as plt + if 'Rqz.ort' in self.output_config.outputFormats: + plt.figure(num=99) + plt.legend() + plt.show() + def read_file_block(self, i, short_notation): logging.warning('reading input:') file_list = self.path_resolver.resolve(short_notation) @@ -149,7 +156,7 @@ class AmorReduction: if self.reduction_config.timeSlize: if i>0: - logging.warning(" time slizing should only be used for on set of datafiles, check parameters") + logging.warning(" time slizing should only be used for one set of datafiles, check parameters") self.analyze_timeslices(i) else: self.analyze_unsliced(i) @@ -188,6 +195,12 @@ class AmorReduction: orso_data = fileio.OrsoDataset(headerRqz, result.data) self.last_result = result self.datasetsRqz.append(orso_data) + + if self.output_config.plot: + import matplotlib.pyplot as plt + # plot all reflectivity results in same graph + plt.figure(num=99) + result.plot(label=f'{self.reduction_config.fileIdentifier[i]}') if 'Rlt.ort' in self.output_config.outputFormats: columns = [ fileio.Column('Qz', '1/angstrom', 'normal momentum transfer'), @@ -230,6 +243,12 @@ class AmorReduction: self.datasetsRlt.append(orso_data) j += 1 + if self.output_config.plot: + import matplotlib.pyplot as plt + plt.figure() + proj.plot(colorbar=True, cmap=str(self.output_config.plot_colormap)) + plt.title(f'{self.reduction_config.fileIdentifier[i]}') + def analyze_timeslices(self, i): wallTime_e = np.float64(self.dataset.data.events.wallTime)/1e9 pulseTimeS = np.float64(self.dataset.data.pulses.time)/1e9 @@ -283,6 +302,13 @@ class AmorReduction: headerRqz.data_set = f'{i}_{ti}: time = {time:8.1f} s to {time+interval:8.1f} s' orso_data = fileio.OrsoDataset(headerRqz, result.data_for_time(time)) self.datasetsRqz.append(orso_data) + + if self.output_config.plot: + import matplotlib.pyplot as plt + # plot all reflectivity results in same graph + plt.figure(num=99) + result.plot(label=f'{self.reduction_config.fileIdentifier[i]} @ {time:.1f}s') + self.last_result = result # reset normal logging behavior #logging.StreamHandler.terminator = "\n" diff --git a/tests/test_full_analysis.py b/tests/test_full_analysis.py index 2105814..34a8321 100644 --- a/tests/test_full_analysis.py +++ b/tests/test_full_analysis.py @@ -7,7 +7,7 @@ from eos import options, reduction, logconfig logconfig.setup_logging() logconfig.update_loglevel(1) -# TODO: add test for new features like proton charge normalization +# TODO: add unit tests for individual parts of reduction class FullAmorTest(TestCase): @classmethod