Implement spectrum copy(), *, *= methods, --scale command-line option

CHANGELOG.rst

@@ -23,6 +23,20 @@
     interpolation method already available for adaptive broadening of
     DOS.
 
+  - Added features to Spectrum classes
+
+    - Added ``copy()`` methods returning an independent duplicate of data
+
+    - Added ``__mul__`` and ``__imul__`` methods to Spectrum
+      classes. This allows results to be conveniently scaled with
+      infix notation ``*`` or ``*=``
+
+  - Added `--scale` parameter to ``euphonic-dos``,
+    ``euphonic-intensity-map``, ``euphonic-powder-map`` to allow
+    arbitrary scaling of results from command-line. (e.g. for
+    comparison with experiment, or changing DOS normalisation from 1
+    to 3N.)
+
 - Bug Fixes:
 
   - Changed the masking logic for kinematic constraints: instead of

euphonic/cli/dos.py

@@ -125,6 +125,9 @@ def energy_broadening_func(x):
     plot_label_kwargs = _plot_label_kwargs(
         args, default_xlabel=f"Energy / {dos.x_data.units:~P}")
 
+    if args.scale is not None:
+        dos *= args.scale
+
     if args.save_json:
         dos.to_json_file(args.save_json)
     style = _compose_style(user_args=args, base=[base_style])
@@ -137,7 +140,8 @@ def get_parser() -> ArgumentParser:
     parser, _ = _get_cli_parser(features={'read-fc', 'read-modes', 'mp-grid',
                                           'plotting', 'ebins',
                                           'adaptive-broadening',
-                                          'pdos-weighting'})
+                                          'pdos-weighting',
+                                          'scaling'})
     parser.description = (
         'Plots a DOS from the file provided. If a force '
         'constants file is provided, a DOS is generated on the Monkhorst-Pack '

euphonic/cli/intensity_map.py

@@ -85,6 +85,9 @@ def main(params: Optional[List[str]] = None) -> None:
     if x_tick_labels:
         spectrum.x_tick_labels = x_tick_labels
 
+    if args.scale is not None:
+        spectrum *= args.scale
+
     spectra = spectrum.split(**split_args)  # type: List[Spectrum2D]
     if len(spectra) > 1:
         print(f"Found {len(spectra)} regions in q-point path")
@@ -104,7 +107,7 @@ def main(params: Optional[List[str]] = None) -> None:
 def get_parser() -> ArgumentParser:
     parser, sections = _get_cli_parser(
         features={'read-fc', 'read-modes', 'q-e', 'map', 'btol', 'ebins',
-                  'ins-weighting', 'plotting'})
+                  'ins-weighting', 'plotting', 'scaling'})
     parser.description = (
         'Plots a 2D intensity map from the file provided. If a force '
         'constants file is provided, a band structure path is '

euphonic/cli/powder_map.py

@@ -39,7 +39,7 @@ def get_parser() -> ArgumentParser:
     parser, sections = _get_cli_parser(
         features={'read-fc', 'pdos-weighting', 'ins-weighting',
                   'powder', 'plotting', 'ebins', 'q-e', 'map',
-                  'brille', 'kinematic'})
+                  'brille', 'kinematic', 'scaling'})
 
     sections['q'].description = (
         '"GRID" options relate to Monkhorst-Pack sampling for the '
@@ -237,6 +237,9 @@ def main(params: Optional[List[str]] = None) -> None:
         spectrum = apply_kinematic_constraints(
             spectrum, e_i=e_i, e_f=e_f, angle_range=args.angle_range)
 
+    if args.scale is not None:
+        spectrum *= args.scale
+
     print(f"Plotting figure: max intensity "
           f"{np.nanmax(spectrum.z_data.magnitude) * spectrum.z_data.units:~P}")
     plot_label_kwargs = _plot_label_kwargs(

euphonic/cli/utils.py

@@ -622,6 +622,10 @@ def __call__(self, parser, args, values, option_string=None):
             '--grid-spacing', type=float, default=0.1, dest='grid_spacing',
             help=('q-point spacing of Monkhorst-Pack grid.'))
 
+    if 'scaling' in features:
+        sections['property'].add_argument(
+            '--scale', type=float, help='Intensity scale factor', default=None)
+
     if 'powder' in features:
         _sampling_choices = {'golden', 'sphere-projected-grid',
                              'spherical-polar-grid',

euphonic/spectra.py

@@ -4,7 +4,7 @@
 import itertools
 import math
 import json
-from numbers import Integral
+from numbers import Integral, Real
 from typing import (Any, Callable, Dict, List, Optional, overload,
                     Sequence, Tuple, TypeVar, Union, Type)
 import warnings
@@ -59,6 +59,22 @@
         self.y_data_unit = str(value.units)
         self._y_data = value.to(self._internal_y_data_unit).magnitude
 
+    def __imul__(self: T, other: Real) -> T:
+        """Scale spectral data in-place"""
+        self._y_data *= other
+        return self
+
+    def __mul__(self: T, other: Real) -> T:
+        """Get a new spectrum with scaled data"""
+        new_spec = self.copy()
+        new_spec *= other
+        return new_spec
+
+    @abstractmethod
+    def copy(self: T) -> T:
+        """Get an independent copy of spectrum"""
+        ...
+
     @property
     def x_tick_labels(self) -> List[Tuple[int, str]]:
         return self._x_tick_labels
@@ -458,6 +474,13 @@
                            metadata=self.metadata)
                 for x0, x1 in ranges]
 
+    def copy(self: T) -> T:
+        """Get an independent copy of spectrum"""
+        return type(self)(np.copy(self.x_data),
+                          np.copy(self.y_data),
+                          x_tick_labels=copy.copy(self.x_tick_labels),
+                          metadata=copy.deepcopy(self.metadata))
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Convert to a dictionary. See Spectrum1D.from_dict for details on
@@ -636,11 +659,9 @@
         else:
             raise TypeError("x_width must be a Quantity or Callable")
 
-        return type(self)(
-            np.copy(self.x_data.magnitude)*ureg(self.x_data_unit),
-            y_broadened,
-            copy.copy((self.x_tick_labels)),
-            copy.copy(self.metadata))
+        new_spectrum = self.copy()
+        new_spectrum.y_data = y_broadened
+        return new_spectrum
 
 
 LineData = Sequence[Dict[str, Union[str, int]]]
@@ -857,7 +878,7 @@
         # Put all other per-spectrum metadata in line_data
         line_data = []
         for i, metadata in enumerate(all_metadata):
-            sdata = copy.copy(metadata)
+            sdata = copy.deepcopy(metadata)
             for key in combined_metadata.keys():
                 sdata.pop(key)
             line_data.append(sdata)
@@ -908,6 +929,10 @@
             line_data_vals[i] = tuple([data[key] for key in line_data_keys])
         return line_data_vals
 
+    def copy(self: T) -> T:
+        """Get an independent copy of spectrum"""
+        return Spectrum1D.copy(self)
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Convert to a dictionary consistent with from_dict()
@@ -1022,13 +1047,14 @@
                 ) -> T:  # noqa: F811
         ...
 
-    def broaden(self: T, x_width: Union[Quantity, CallableQuantity],
+    def broaden(self: T,
+                x_width: Union[Quantity, CallableQuantity],
                 shape: str = 'gauss',
                 method: Optional[str] = None,
                 width_lower_limit: Quantity = None,
                 width_convention: str = 'FWHM',
                 width_interpolation_error: float = 0.01,
-                width_fit: str= 'cheby-log'
+                width_fit: str = 'cheby-log'
                 ) -> T:  # noqa: F811
         """
         Individually broaden each line in y_data, returning a new
@@ -1084,11 +1110,10 @@
                 y_broadened[i] = self._broaden_data(
                     yi, x_centres, x_width_calc, shape=shape,
                     method=method)
-            return Spectrum1DCollection(
-                np.copy(self.x_data.magnitude)*ureg(self.x_data_unit),
-                y_broadened*ureg(self.y_data_unit),
-                copy.copy((self.x_tick_labels)),
-                copy.deepcopy(self.metadata))
+
+            new_spectrum = self.copy()
+            new_spectrum.y_data = y_broadened * ureg(self.y_data_unit)
+            return new_spectrum
 
         elif isinstance(x_width, Callable):
             return type(self).from_spectra([
@@ -1141,9 +1166,11 @@
         new_y_data = new_y_data*ureg(self._internal_y_data_unit).to(
             self.y_data_unit)
 
-        return Spectrum1DCollection(self.x_data, new_y_data,
-                                    x_tick_labels=self.x_tick_labels,
-                                    metadata=group_metadata)
+        new_data = self.copy()
+        new_data.y_data = new_y_data
+        new_data.metadata = group_metadata
+
+        return new_data
 
     def sum(self) -> Spectrum1D:
         """
@@ -1161,9 +1188,10 @@
         metadata.update(self._combine_line_metadata())
         summed_y_data = np.sum(self._y_data, axis=0)*ureg(
             self._internal_y_data_unit).to(self.y_data_unit)
-        return Spectrum1D(self.x_data, summed_y_data,
-                          x_tick_labels=self.x_tick_labels,
-                          metadata=metadata)
+        return Spectrum1D(np.copy(self.x_data),
+                          summed_y_data,
+                          x_tick_labels=copy.copy(self.x_tick_labels),
+                          metadata=copy.deepcopy(metadata))
 
     def select(self, **select_key_values: Union[
             str, int, Sequence[str], Sequence[int]]) -> T:
@@ -1293,6 +1321,11 @@
         self.z_data_unit = str(value.units)
         self._z_data = value.to(self._internal_z_data_unit).magnitude
 
+    def __imul__(self: T, other: Real) -> T:
+        """Scale spectral data in-place"""
+        self.z_data = self.z_data * other
+        return self
+
     def __setattr__(self, name: str, value: Any) -> None:
         _check_unit_conversion(self, name, value,
                                ['z_data_unit'])
@@ -1403,7 +1436,7 @@
                 np.copy(self.y_data.magnitude)*ureg(self.y_data_unit),
                 z_broadened*ureg(self.z_data_unit),
                 copy.copy(self.x_tick_labels),
-                copy.copy(self.metadata))
+                copy.deepcopy(self.metadata))
         else:
             spectrum = self
 
@@ -1475,6 +1508,14 @@
             copy.copy(spectrum.x_tick_labels),
             copy.copy(spectrum.metadata))
 
+    def copy(self: T) -> T:
+        """Get an independent copy of spectrum"""
+        return type(self)(np.copy(self.x_data),
+                          np.copy(self.y_data),
+                          np.copy(self.z_data),
+                          copy.copy(self.x_tick_labels),
+                          copy.deepcopy(self.metadata))
+
     def get_bin_edges(self, bin_ax: str = 'x') -> Quantity:
         """
         Get bin edges for the axis specified by bin_ax. If the size of
@@ -1695,21 +1736,15 @@
                                                            float('-Inf')]
     q_bounds = q_bounds.real
 
-    new_z_data = np.copy(spectrum.z_data.magnitude)
-
     mask = np.logical_or((spectrum.get_bin_edges(bin_ax='x')[1:, np.newaxis]
                           < q_bounds[0][np.newaxis, :]),
                          (spectrum.get_bin_edges(bin_ax='x')[:-1, np.newaxis]
                           > q_bounds[-1][np.newaxis, :]))
 
-    new_z_data[mask] = float('nan')
+    new_spectrum = spectrum.copy()
+    new_spectrum._z_data[mask] = float('nan')
 
-    return Spectrum2D(
-        np.copy(spectrum.x_data.magnitude) * ureg(spectrum.x_data_unit),
-        np.copy(spectrum.y_data.magnitude) * ureg(spectrum.y_data_unit),
-        new_z_data * ureg(spectrum.z_data_unit),
-        copy.copy(spectrum.x_tick_labels),
-        copy.deepcopy(spectrum.metadata))
+    return new_spectrum
 
 
 def _get_cos_range(angle_range: Tuple[float]) -> Tuple[float]: