1"""
  2# Description
  3
  4This module contains common classes used to load and manipulate spectral data.
  5Any class can be instantiated directly as `sah.Class()`.
  6
  7
  8# Index
  9
 10| | |
 11| --- | --- |
 12| `Spectra`  | Used to load and process spectral data |
 13| `Plotting` | Stores plotting options, used in `Spectra.plotting` |
 14| `Material` | Used to store and calculate material parameters, such as molar masses and neutron cross sections |
 15
 16
 17# Examples
 18
 19To load two INS spectra CSV files with cm$^{-1}$ as input units,
 20converting them to meV units, and finally plotting them:
 21```python
 22import sah
 23ins = sah.Spectra(
 24    type     = 'ins',
 25    files    = ['example_1.csv', 'example_2.csv'],
 26    units_in = 'cm-1',
 27    units    = 'meV',
 28    )
 29sah.plot(ins)
 30```
 31
 32Check more use examples in the [`examples/`](https://github.com/pablogila/sah/tree/main/examples) folder.
 33
 34---
 35"""
 36
 37
 38import numpy as np
 39import pandas as pd
 40from copy import deepcopy
 41import os
 42import aton
 43import periodictable
 44import scipy
 45
 46
 47# Common conversion factors
 48cm1_to_meV = (scipy.constants.h * scipy.constants.c * 100 / scipy.constants.e) * 1000
 49meV_to_cm1 = 1 / cm1_to_meV
 50
 51
 52class Plotting:
 53    """Stores plotting options, read by `sah.plot`"""
 54    def __init__(
 55            self,
 56            title:str=None,
 57            xlim=None,
 58            ylim=None,
 59            margins:list=[0,0],
 60            offset=True,
 61            scaling:float=1.0,
 62            vline:list=None,
 63            vline_error:list=None,
 64            figsize:tuple=None,
 65            log_xscale:bool=False,
 66            show_yticks:bool=False,
 67            xlabel:str=None,
 68            ylabel:str=None,
 69            legend=None,
 70            legend_title:str=None,
 71            legend_size='medium',
 72            legend_loc='best',
 73            viridis:bool=False,
 74            save_as:str=None,
 75        ):
 76        """Default values can be overwritten when initializing the Plotting object."""
 77        self.title = title
 78        """Title of the plot. Set it to an empty string to remove the title."""
 79        self.xlim = self._set_limits(xlim)
 80        """List with the x-limits of the plot, as in `[xlim_low, xlim_top]`."""
 81        self.ylim = self._set_limits(ylim)
 82        """List with the y-limits of the plot, as in `[ylim_low, ylim_top]`."""
 83        self.margins = self._set_limits(margins)
 84        """List with additional margins at the bottom and top of the plot, as in `[low_margin, top_margin]`."""
 85        self.offset = offset
 86        """If `True`, the plots will be separated automatically.
 87
 88        It can be set to a float, to equally offset the plots by a given value.
 89        """
 90        self.scaling = scaling
 91        "Scaling factor"
 92        if vline is not None and not isinstance(vline, list):
 93            vline = [vline]
 94        self.vline = vline
 95        """Vertical line/s to plot. Can be an int or float with the x-position, or a list with several ones."""
 96        if vline_error is not None and not isinstance(vline_error, list):
 97            vline_error = [vline_error]
 98        self.vline_error = vline_error
 99        """Plot a shaded area of the specified width around the vertical lines specified at `vline`.
100
101        It can be an array of the same length as `vline`, or a single value to be applied to all.
102        """
103        self.figsize = figsize
104        """Tuple with the figure size, as in matplotlib."""
105        self.log_xscale = log_xscale
106        """If true, plot the x-axis in logarithmic scale."""
107        self.show_yticks = show_yticks
108        """Show or not the yticks on the plot."""
109        self.xlabel = xlabel
110        """Custom label of the x-axis.
111
112        If `None`, the default label will be used.
113        Set to `''` to remove the label of the horizontal axis.
114        """
115        self.ylabel = ylabel
116        """Custom label of the y-axis.
117        
118        If `None`, the default label will be used.
119        Set to `''` to remove the label of the vertical axis.
120        """
121        if not isinstance(legend, list) and legend is not None and legend != False:
122            legend = [legend]
123        self.legend = legend
124        """Legend of the plot.
125
126        If `None`, the filenames will be used as legend.
127        Can be a bool to show or hide the plot legend.
128        It can also be an array containing the strings to display;
129        in that case, elements set to `False` will not be displayed.
130        """
131        self.legend_title = legend_title
132        """Title of the legend, defaults to `None`."""
133        self.legend_size = legend_size
134        """Size of the legend, as in matplotlib. Defaults to `'medium'`."""
135        self.legend_loc = legend_loc
136        """Location of the legend, as in matplotlib. Defaults to `'best'`."""
137        self.viridis: bool = viridis
138        """Use the Viridis colormap for the plot. Defaults to `False`."""
139        self.save_as = save_as
140        """Filename to save the plot. None by default."""
141
142    def _set_limits(self, limits) -> list:
143        """Set the x and y limits of the plot."""
144        if limits is None:
145            return [None, None]
146        if isinstance(limits, tuple):
147            limits = list(limits)
148        if isinstance(limits, list):
149            if len(limits) == 0:
150                return [None, None]
151            if len(limits) == 1:
152                return [None, limits[0]]
153            if len(limits) == 2:
154                return limits
155            else:
156                return limits[:2]
157        if isinstance(limits, int) or isinstance(limits, float):
158            return [None, limits]
159        else:
160            raise ValueError(f"Unknown plotting limits: Must be specified as a list of two elements, as [low_limit, high_limit]. Got: {limits}")
161
162
163class Spectra:
164    """Spectra object. Used to load and process spectral data.
165
166    Most functions in the `sah` module receive this object as input.
167    """
168    def __init__(
169            self,
170            type:str=None,
171            comment:str=None,
172            files=None,
173            dfs=None,
174            units=None,
175            units_in=None,
176            plotting:Plotting=Plotting(),
177        ):
178        """All values can be set when initializing the Spectra object."""
179        self.type = None
180        """Type of the spectra: `'INS'`, `'ATR'`, or `'RAMAN'`."""
181        self.comment = comment
182        """Custom comment. If `Plotting.title` is None,  it will be the title of the plot."""
183        self.files = None
184        """List containing the files with the spectral data.
185
186        Loaded automatically to `dfs` with Pandas at initialization.
187        In order for Pandas to read the files properly, note that the column lines must start by `#`.
188        Any additional line that is not data must be removed or commented with `#`.
189        CSV files must be formatted with the first column as the energy or energy transfer,
190        and the second column with the intensity or absorbance, depending on the case.
191        An additional third `'Error'` column can be used.
192        """
193        self.dfs = None
194        """List containing the pandas dataframes with the spectral data.
195        
196        Loaded automatically from `files` at initialization.
197        """
198        self.units = None
199        """Target units of the spectral data.
200        
201        Can be `'meV'` or `'cm-1'`."""
202        self.units_in = None
203        """Input units of the spectral data, used in the input CSV files.
204        
205        Can be `'meV'` or `'cm-1'`.
206        If the input CSV files have different units,
207        it can also be set as a list of the same length of the number of input files,
208        eg. `['meV', 'cm-1', 'cm-1']`.
209        """
210        self.plotting = plotting
211        """`Plotting` object, used to set the plotting options."""
212
213        self = self._set_type(type)
214        self = self._set_dataframes(files, dfs)
215        self = self.set_units(units, units_in)
216
217    def _set_type(self, type):
218        """Set and normalize the type of the spectra: `INS`, `ATR`, or `RAMAN`."""
219        if type.lower() in aton.alias.experiments['ins']:
220            self.type = 'ins'
221        elif type.lower() in aton.alias.experiments['atr']:
222            self.type = 'atr'
223        elif type.lower() in aton.alias.experiments['raman']:
224            self.type = 'raman'
225        else:
226            self.type = type.lower()
227        return self
228
229    def _set_dataframes(self, files, dfs):
230        '''Set the dfs list of dataframes, from the given files or dfs.'''
231        if isinstance(files, list):
232            self.files = files
233        elif isinstance(files, str):
234            self.files = [files]
235        else:
236            self.files = []
237
238        if isinstance(dfs, pd.DataFrame):
239            self.dfs = [dfs]
240        elif isinstance(dfs, list) and isinstance(dfs[0], pd.DataFrame):
241            self.dfs = dfs
242        else:
243            self.dfs = [self._read_dataframe(filename) for filename in self.files]
244        return self
245
246    def _read_dataframe(self, filename):
247        """Read a dataframe from a file."""
248        root = os.getcwd()
249        file_path = os.path.join(root, filename)
250        df = pd.read_csv(file_path, comment='#', sep=r',|;|\s+', engine='python', header=None)
251        # Remove any empty columns
252        df = df.dropna(axis=1, how='all')
253        df = df.sort_values(by=df.columns[0]) # Sort the data by energy
254
255        print(f'\nNew dataframe from {filename}')
256        print(df.head(),'\n')
257        return df
258
259    def set_units(
260            self,
261            units,
262            units_in=None,
263            default_unit='cm-1',
264            ):
265        """Method to change between spectral units. ALWAYS use this method to do that.
266
267        For example, to change to meV from cm-1:
268        ```python
269        Spectra.set_units('meV', 'cm-1')
270        ```
271        """
272        mev = 'meV'
273        cm = 'cm-1'
274        unit_format={
275                mev: aton.alias.units['meV'],
276                cm: aton.alias.units['cm1'] + aton.alias.units['cm'],
277            }
278        if self.units is not None:
279            units_in = deepcopy(self.units)
280            self.units = units
281        elif units is not None:
282            units_in = units_in
283            self.units = deepcopy(units)
284        elif units is None and units_in is None:
285            units_in = None
286            self.units = default_unit
287        elif units is None and units_in is not None:
288            units_in = None
289            self.units = deepcopy(units_in)
290        if isinstance(units_in, list):
291            for i, unit_in in enumerate(units_in):
292                for key, value in unit_format.items():
293                    if unit_in in value:
294                        units_in[i] = key
295                        break
296            if len(units_in) == 1:
297                units_in = units_in * len(self.files)
298            elif len(units_in) != len(self.files):
299                raise ValueError("units_in must be a list of the same length as files.")
300        if isinstance(units_in, str):
301            for key, value in unit_format.items():
302                if units_in in value:
303                    units_in = key
304                    break
305            units_in = [units_in] * len(self.files)
306        if isinstance(self.units, list):
307            for i, unit in enumerate(self.units):
308                for key, value in unit_format.items():
309                    if unit in value:
310                        self.units[i] = key
311                        break
312            if len(self.units) == 1:
313                self.units = self.units * len(self.files)
314            elif len(self.units) != len(self.files):
315                raise ValueError("units_in must be a list of the same length as files.")
316        if isinstance(self.units, str):
317            for key, value in unit_format.items():
318                if self.units in value:
319                    self.units = key
320                    break
321            self.units = [self.units] * len(self.files)
322        if units_in is None:
323            return self
324        # Otherwise, convert the dfs
325        if len(self.units) != len(units_in):
326            raise ValueError("Units len mismatching.")
327        for i, unit in enumerate(self.units):
328            if unit == units_in[i]:
329                continue
330            if unit == mev and units_in[i] == cm: 
331                self.dfs[i][self.dfs[i].columns[0]] = self.dfs[i][self.dfs[i].columns[0]] * cm1_to_meV
332        for i, df in enumerate(self.dfs):
333            if self.units[i] == mev:
334                E_units = 'meV'
335            elif self.units[i] == cm:
336                E_units = 'cm-1'
337            else:
338                E_units = self.units[i]
339            if self.type == 'ins':
340                if self.dfs[i].shape[1] == 3:
341                    self.dfs[i].columns = [f'Energy transfer / {E_units}', 'S(Q,E)', 'Error']
342                elif self.dfs[i].shape[1] == 2:
343                    self.dfs[i].columns = [f'Energy transfer / {E_units}', 'S(Q,E)']
344            elif self.type == 'atr':
345                if self.dfs[i].shape[1] == 3:
346                    self.dfs[i].columns = [f'Wavenumber / {E_units}', 'Absorbance', 'Error']
347                elif self.dfs[i].shape[1] == 2:
348                    self.dfs[i].columns = [f'Wavenumber / {E_units}', 'Absorbance']
349            elif self.type == 'raman':
350                if self.dfs[i].shape[1] == 3:
351                    self.dfs[i].columns = [f'Raman shift / {E_units}', 'Counts', 'Error']
352                elif self.dfs[i].shape[1] == 2:
353                    self.dfs[i].columns = [f'Raman shift / {E_units}', 'Counts']
354        return self
355
356
357class Material:
358    """Material class.
359
360    Used to calculate molar masses and cross sections,
361    and to pass data to different analysis functions
362    such as `sah.deuterium.impulse_approx().`
363    """
364    def __init__(
365            self,
366            elements:dict,
367            name:str=None,
368            grams:float=None,
369            grams_error:float=None,
370            mols:float=None,
371            mols_error:float=None,
372            molar_mass:float=None,
373            cross_section:float=None,
374            peaks:dict=None,
375        ):
376        """
377        All values can be set when initializing the Material object.
378        However, it is recommended to only set the elements and the grams,
379        and optionally the name, and calculate the rest with `Material.set()`.
380        """
381        self.elements = elements
382        """Dict of atoms in the material, as in `{'C':1, 'N':1, 'H': 3, 'D':3}`.
383
384        Isotopes can be expressed as 'H2', 'He4', etc. with the atom symbol + isotope mass number.
385        """
386        self.name = name
387        """String with the name of the material."""
388        self.grams = grams
389        """Mass, in grams."""
390        self.grams_error = grams_error
391        """Error of the measured mass in grams.
392
393        Set automatically with `Material.set()`.
394        """
395        self.mols = mols
396        """Number of moles.
397
398        Set automatically with `Material.set()`.
399        """
400        self.mols_error = mols_error
401        """Error of the number of moles.
402
403        Set automatically with `Material.set()`.
404        """
405        self.molar_mass = molar_mass
406        """Molar mass of the material, in mol/g.
407
408        Calculated automatically with `Material.set()`.
409        """
410        self.cross_section = cross_section
411        """Neutron total bound scattering cross section, in barns.
412
413        Calculated automatically with `Material.set()`.
414        """
415        self.peaks = peaks
416        """Dict with interesting peaks that you might want to store for later use."""
417
418    def _set_grams_error(self):
419        """Set the error in grams, based on the number of decimal places."""
420        if self.grams is None:
421            return
422        decimal_accuracy = len(str(self.grams).split('.')[1])
423        # Calculate the error in grams
424        self.grams_error = 10**(-decimal_accuracy)
425
426    def _set_mass(self):
427        """Set the molar mass of the material.
428
429        If `Material.grams` is provided, the number of moles will be
430        calculated and overwritten. Isotopes can be used as 'element + A',
431        eg. `'He4'`. This gets splitted with `aton.txt.extract.isotope()`.
432        """
433        material_grams_per_mol = 0.0
434        for key in self.elements:
435            try:
436                material_grams_per_mol += self.elements[key] * periodictable.elements.symbol(key).mass
437            except KeyError:  # Split the atomic flag as H2, etc
438                element, isotope = aton.txt.extract.isotope(key)
439                isotope_name = isotope+'-'+element  # Periodictable format
440                material_grams_per_mol += self.elements[key] * periodictable.elements.isotope(isotope_name).mass
441        self.molar_mass = material_grams_per_mol
442        if self.grams is not None:
443            self._set_grams_error()
444            self.mols = self.grams / material_grams_per_mol
445            self.mols_error = self.mols * np.sqrt((self.grams_error / self.grams)**2)
446    
447    def _set_cross_section(self):
448        """Set the cross section of the material, based on the `elements` dict.
449
450        If an isotope is used, eg. `'He4'`, it splits the name with `aton.txt.extract.isotope()`.
451        """
452        total_cross_section = 0.0
453        for key in self.elements:
454            try:
455                total_cross_section += self.elements[key] * periodictable.elements.symbol(key).neutron.total
456            except KeyError: # Split the atomic flag as H2, etc
457                element, isotope_index = aton.txt.extract.isotope(key)
458                isotope_name = isotope+'-'+element  # Periodictable format
459                total_cross_section += self.elements[key] * periodictable.elements.isotope(isotope_name).neutron.total
460        self.cross_section = total_cross_section
461
462    def set(self):
463        """Set the molar mass, cross section and errors of the material."""
464        self._set_mass()
465        self._set_cross_section()
466
467    def print(self):
468        """Print a summary with the material information."""
469        print('\nMATERIAL')
470        if self.name is not None:
471            print(f'Name: {self.name}')
472        if self.grams is not None and self.grams_error is not None:
473            print(f'Grams: {self.grams} +- {self.grams_error} g')
474        elif self.grams is not None:
475            print(f'Grams: {self.grams} g')
476        if self.mols is not None and self.mols_error is not None:
477            print(f'Moles: {self.mols} +- {self.mols_error} mol')
478        elif self.mols is not None:
479            print(f'Moles: {self.mols} mol')
480        if self.molar_mass is not None:
481            print(f'Molar mass: {self.molar_mass} g/mol')
482        if self.cross_section is not None:
483            print(f'Cross section: {self.cross_section} barns')
484        if self.elements is not None:
485            print(f'Elements: {self.elements}')
486        print('')