rubin-sim 2.6.1.dev2__py3-none-any.whl

rubin_sim/maf/slicers/movie_slicer.py

@@ -0,0 +1,288 @@
+# cumulative one dimensional movie slicer
+__all__ = ("MovieSlicer",)
+
+import os
+import subprocess
+import warnings
+from functools import wraps
+from subprocess import CalledProcessError
+
+import numpy as np
+
+from rubin_sim.maf.stackers import ColInfo
+from rubin_sim.maf.utils import optimal_bins
+
+from .base_slicer import BaseSlicer
+
+
+class MovieSlicer(BaseSlicer):
+    def __init__(
+        self,
+        slice_col_name=None,
+        slice_col_units=None,
+        bins=None,
+        bin_min=None,
+        bin_max=None,
+        bin_size=None,
+        verbose=True,
+        badval=0,
+        cumulative=True,
+        force_no_ffmpeg=False,
+    ):
+        """The 'movieSlicer' acts similarly to the OneDSlicer
+        (slices on one data column).
+        However, the data slices from the movieSlicer are intended to be
+        fed to another slicer, which then
+        (together with a set of Metrics) calculates metric values +
+        plots at each slice created by the movieSlicer.
+        The job of the movieSlicer is to track those slices and
+        put them together into a movie.
+
+        'slice_col_name' is the name of the data column to use for slicing.
+        'slice_col_units' lets the user set the units (for plotting purposes)
+        of the slice column.
+        'bins' can be a numpy array with the binpoints for sliceCol
+        or a single integer value
+        (if a single value, this will be used as the number of bins,
+        together with data min/max or bin_min/max),
+        as in numpy's histogram function.
+        If 'bin_size' is used, this will override the bins value and
+        will be used together with the data min/max
+        or bin_min/Max to set the binpoint values.
+
+        Bins work like numpy histogram bins: the last 'bin' value is end
+        value of last bin; all bins except for last bin are half-open
+        ([a, b>), the last one is ([a, b]).
+
+        The movieSlicer stitches individual frames together into a movie
+        using ffmpeg. Thus, on instantiation it checks that ffmpeg is
+        available and will raise and exception if not.
+        This behavior can be overriden using force_no_ffmpeg = True
+        (in order to create a movie later perhaps).
+        """
+        # Check for ffmpeg.
+        if not force_no_ffmpeg:
+            try:
+                fnull = open(os.devnull, "w")
+                subprocess.check_call(["which", "ffmpeg"], stdout=fnull)
+            except CalledProcessError:
+                raise Exception(
+                    "Could not find ffmpeg on the system, so will not be able to create movie."
+                    " Use force_no_ffmpeg=True to override this error and create individual images."
+                )
+        super().__init__(verbose=verbose, badval=badval)
+        self.slice_col_name = slice_col_name
+        self.columns_needed = [slice_col_name]
+        self.bins = bins
+        self.bin_min = bin_min
+        self.bin_max = bin_max
+        self.bin_size = bin_size
+        self.cumulative = cumulative
+        if slice_col_units is None:
+            co = ColInfo()
+            self.slice_col_units = co.get_units(self.slice_col_name)
+        self.slicer_init = {
+            "slice_col_name": self.slice_col_name,
+            "slice_col_units": slice_col_units,
+            "badval": badval,
+        }
+
+    def setup_slicer(self, sim_data, maps=None):
+        """
+        Set up bins in slicer.
+        """
+        if self.slice_col_name is None:
+            raise Exception("slice_col_name was not defined when slicer instantiated.")
+        slice_col = sim_data[self.slice_col_name]
+        # Set bin min/max values.
+        if self.bin_min is None:
+            self.bin_min = slice_col.min()
+        if self.bin_max is None:
+            self.bin_max = slice_col.max()
+        # Give warning if bin_min = bin_max, and do something at
+        # least slightly reasonable.
+        if self.bin_min == self.bin_max:
+            warnings.warn(
+                "bin_min = bin_max (maybe your data is single-valued?). "
+                "Increasing bin_max by 1 (or 2*bin_size, if bin_size set)."
+            )
+            if self.bin_size is not None:
+                self.bin_max = self.bin_max + 2 * self.bin_size
+            else:
+                self.bin_max = self.bin_max + 1
+        # Set bins.
+        # Using bin_size.
+        if self.bin_size is not None:
+            if self.bins is not None:
+                warnings.warn(
+                    "Both bin_size and bins have been set; Using bin_size %f only." % (self.bin_size)
+                )
+            self.bins = np.arange(
+                self.bin_min,
+                self.bin_max + self.bin_size / 2.0,
+                float(self.bin_size),
+                "float",
+            )
+        # Using bins value.
+        else:
+            # Bins was a sequence (np array or list)
+            if hasattr(self.bins, "__iter__"):
+                self.bins = np.sort(self.bins)
+            # Or bins was a single value.
+            else:
+                if self.bins is None:
+                    self.bins = optimal_bins(slice_col, self.bin_min, self.bin_max)
+                nbins = int(self.bins)
+                self.bin_size = (self.bin_max - self.bin_min) / float(nbins)
+                self.bins = np.arange(
+                    self.bin_min,
+                    self.bin_max + self.bin_size / 2.0,
+                    self.bin_size,
+                    "float",
+                )
+        # Set shape / nbins to be one less than # of bins because
+        # last binvalue is RH edge only
+        self.nslice = len(self.bins) - 1
+        self.shape = self.nslice
+        # Set slice_point metadata.
+        self.slice_points["sid"] = np.arange(self.nslice)
+        self.slice_points["bins"] = self.bins
+        # Add metadata from maps.
+        self._run_maps(maps)
+        # Set up data slicing.
+        self.sim_idxs = np.argsort(sim_data[self.slice_col_name])
+        simFieldsSorted = np.sort(sim_data[self.slice_col_name])
+        # "left" values are location where simdata == bin value
+        # Note that this setup will place any the visits beyond
+        # the last bin into the last bin.
+        self.left = np.searchsorted(simFieldsSorted, self.bins[:-1], "left")
+        self.left = np.concatenate(
+            (
+                self.left,
+                np.array(
+                    [
+                        len(self.sim_idxs),
+                    ]
+                ),
+            )
+        )
+        # Set up _slice_sim_data method for this class.
+        if self.cumulative:
+
+            @wraps(self._slice_sim_data)
+            def _slice_sim_data(islice):
+                """
+                Slice sim_data on oneD sliceCol, to
+                return relevant indexes for slice_point.
+                """
+                # this is the important part.
+                # The ids here define the pieces of data that get
+                # passed on to subsequent slicers
+                # cumulative version of 1D slicing
+                idxs = self.sim_idxs[0 : self.left[islice + 1]]
+                return {
+                    "idxs": idxs,
+                    "slice_point": {
+                        "sid": islice,
+                        "bin_left": self.bins[0],
+                        "bin_right": self.bins[islice + 1],
+                    },
+                }
+
+            setattr(self, "_slice_sim_data", _slice_sim_data)
+        else:
+
+            @wraps(self._slice_sim_data)
+            def _slice_sim_data(islice):
+                """
+                Slice sim_data on oneD sliceCol,
+                to return relevant indexes for slice_point.
+                """
+                idxs = self.sim_idxs[self.left[islice] : self.left[islice + 1]]
+                return {
+                    "idxs": idxs,
+                    "slice_point": {
+                        "sid": islice,
+                        "bin_left": self.bins[islice],
+                        "bin_right": self.bins[islice + 1],
+                    },
+                }
+
+            setattr(self, "_slice_sim_data", _slice_sim_data)
+
+    def __eq__(self, other_slicer):
+        """
+        Evaluate if slicers are equivalent.
+        """
+        if isinstance(other_slicer, MovieSlicer):
+            return np.array_equal(other_slicer.slice_points["bins"], self.slice_points["bins"])
+        else:
+            return False
+
+    def make_movie(
+        self,
+        outfileroot,
+        sliceformat,
+        plot_type,
+        fig_format,
+        out_dir="Output",
+        ips=10.0,
+        fps=10.0,
+    ):
+        """Takes in metric and slicer metadata and calls
+        ffmpeg to stitch together output files.
+        """
+        if not os.path.isdir(out_dir):
+            raise Exception("Cannot find output directory %s with movie input files." % (out_dir))
+        # make video
+        # ffmpeg -r 30 -i [image names - FilterColors_%03d_SkyMap.png] -r 30
+        #    -pix_fmt yuv420p -crf 18 -preset slower outfile
+        callList = [
+            "ffmpeg",
+            "-r",
+            str(ips),
+            "-i",
+            os.path.join(
+                out_dir,
+                "%s_%s_%s.%s" % (outfileroot, sliceformat, plot_type, fig_format),
+            ),
+            "-vf",
+            "pad=ceil(iw/2)*2:ceil(ih/2)*2",
+            "-r",
+            str(fps),
+            "-pix_fmt",
+            "yuv420p",
+            "-crf",
+            "18",
+            "-preset",
+            "slower",
+            os.path.join(
+                out_dir,
+                "%s_%s_%s_%s.mp4" % (outfileroot, plot_type, str(ips), str(fps)),
+            ),
+        ]
+        print("Attempting to call ffmpeg with:")
+        print(" ".join(callList))
+        subprocess.check_call(callList)
+        # make thumbnail gif
+        callList = [
+            "ffmpeg",
+            "-i",
+            os.path.join(
+                out_dir,
+                "%s_%s_%s_%s.mp4" % (outfileroot, plot_type, str(ips), str(fps)),
+            ),
+            "-vf",
+            "pad=ceil(iw/2)*2:ceil(ih/2)*2",
+            "-t",
+            str(10),
+            "-r",
+            str(10),
+            os.path.join(
+                out_dir,
+                "%s_%s_%s_%s.gif" % (outfileroot, plot_type, str(ips), str(fps)),
+            ),
+        ]
+        print("converting to animated gif with:")
+        print(" ".join(callList))
+        subprocess.check_call(callList)

rubin_sim/maf/slicers/nd_slicer.py

@@ -0,0 +1,147 @@
+__all__ = ("NDSlicer",)
+
+import itertools
+import warnings
+from functools import wraps
+
+import numpy as np
+
+from rubin_sim.maf.plots.nd_plotters import OneDSubsetData, TwoDSubsetData
+
+from .base_slicer import BaseSlicer
+
+# nd Slicer slices data on N columns in sim_data
+
+
+class NDSlicer(BaseSlicer):
+    """Nd slicer (N dimensions)
+
+    Parameters
+    ==========
+    slice_col_list : `list` of `str`
+        Names of the data columns for slicing (e.g. 'airmass`, etc)
+    bins_list : `int` or `list` of `int` or `np.ndarray`, opt
+        Single integer (for same number of slices in each dimension)
+        or a list of integers (matching
+        slice_col_list) or list of arrays. Default 100, in all dimensions.
+
+    All bins are half-open ([a, b)).
+    """
+
+    def __init__(self, slice_col_list=None, bins_list=100, verbose=True):
+        super().__init__(verbose=verbose)
+        self.slice_col_list = slice_col_list
+        self.columns_needed = self.slice_col_list
+        if self.slice_col_list is not None:
+            self.n_d = len(self.slice_col_list)
+        else:
+            self.n_d = None
+        self.bins_list = bins_list
+        if not (isinstance(bins_list, float) or isinstance(bins_list, int)):
+            if len(self.bins_list) != self.n_d:
+                raise Exception(
+                    "bins_list must be same length as slice_col_names, unless it is a single value"
+                )
+        self.slicer_init = {"slice_col_list": slice_col_list, "bins_list": bins_list}
+        self.plot_funcs = [TwoDSubsetData, OneDSubsetData]
+
+    def setup_slicer(self, sim_data, maps=None):
+        """Set up bins."""
+        # Parse input bins choices.
+        self.bins = []
+        # If we were given a single number for the binsList, convert to list.
+        if isinstance(self.bins_list, float) or isinstance(self.bins_list, int):
+            self.bins_list = [self.bins_list for c in self.slice_col_list]
+        # And then build bins.
+        for bl, col in zip(self.bins_list, self.slice_col_list):
+            if isinstance(bl, float) or isinstance(bl, int):
+                slice_col = sim_data[col]
+                bin_min = slice_col.min()
+                bin_max = slice_col.max()
+                if bin_min == bin_max:
+                    warnings.warn("bin_min=bin_max for column %s: increasing bin_max by 1." % (col))
+                    bin_max = bin_max + 1
+                bin_size = (bin_max - bin_min) / float(bl)
+                bins = np.arange(bin_min, bin_max + bin_size / 2.0, bin_size, "float")
+                self.bins.append(bins)
+            else:
+                self.bins.append(np.sort(bl))
+        self.nslice = (np.array(list(map(len, self.bins))) - 1).prod()
+        # Count how many bins we have total
+        # (not counting last 'RHS' bin values, as in oneDSlicer).
+        self.shape = self.nslice
+        # Set up slice metadata.
+        self.slice_points["sid"] = np.arange(self.nslice)
+        # Including multi-D 'leftmost' bin values
+        bins_for_iteration = []
+        for b in self.bins:
+            bins_for_iteration.append(b[:-1])
+        biniterator = itertools.product(*bins_for_iteration)
+        self.slice_points["bins"] = []
+        for b in biniterator:
+            self.slice_points["bins"].append(b)
+        # and multi-D 'leftmost' bin indexes corresponding to each sid
+        self.slice_points["bin_idxs"] = []
+        bin_ids_for_iteration = []
+        for b in self.bins:
+            bin_ids_for_iteration.append(np.arange(len(b[:-1])))
+        bin_id_iterator = itertools.product(*bin_ids_for_iteration)
+        for bidx in bin_id_iterator:
+            self.slice_points["bin_idxs"].append(bidx)
+        # Add metadata from maps.
+        self._run_maps(maps)
+        # Set up indexing for data slicing.
+        self.sim_idxs = []
+        self.lefts = []
+        for slice_col_name, bins in zip(self.slice_col_list, self.bins):
+            sim_idxs = np.argsort(sim_data[slice_col_name])
+            sim_fields_sorted = np.sort(sim_data[slice_col_name])
+            # "left" values are location where simdata == bin value
+            left = np.searchsorted(sim_fields_sorted, bins[:-1], "left")
+            left = np.concatenate(
+                (
+                    left,
+                    np.array(
+                        [
+                            len(sim_idxs),
+                        ]
+                    ),
+                )
+            )
+            # Add these calculated values into the class lists of
+            # sim_idxs and lefts.
+            self.sim_idxs.append(sim_idxs)
+            self.lefts.append(left)
+
+        @wraps(self._slice_sim_data)
+        def _slice_sim_data(islice):
+            """Slice sim_data to return relevant indexes for slice_point."""
+            # Identify relevant pointings in each dimension.
+            sim_idxs_list = []
+            # Translate islice into indexes in each bin dimension
+            bin_idxs = self.slice_points["bin_idxs"][islice]
+            for d, i in zip(list(range(self.n_d)), bin_idxs):
+                sim_idxs_list.append(set(self.sim_idxs[d][self.lefts[d][i] : self.lefts[d][i + 1]]))
+            idxs = list(set.intersection(*sim_idxs_list))
+            return {
+                "idxs": idxs,
+                "slice_point": {
+                    "sid": islice,
+                    "bin_left": self.slice_points["bins"][islice],
+                    "bin_idxs": self.slice_points["bin_idxs"][islice],
+                },
+            }
+
+        setattr(self, "_slice_sim_data", _slice_sim_data)
+
+    def __eq__(self, other_slicer):
+        """Evaluate if grids are equivalent."""
+        if isinstance(other_slicer, NDSlicer):
+            if other_slicer.n_d != self.n_d:
+                return False
+            for i in range(self.n_d):
+                if not np.array_equal(other_slicer.slice_points["bins"][i], self.slice_points["bins"][i]):
+                    return False
+            return True
+        else:
+            return False

rubin_sim/maf/slicers/one_d_slicer.py

@@ -0,0 +1,210 @@
+__all__ = ("OneDSlicer",)
+
+import warnings
+from functools import wraps
+
+import numpy as np
+
+from rubin_sim.maf.plots.oned_plotters import OneDBinnedData
+from rubin_sim.maf.stackers import ColInfo
+from rubin_sim.maf.utils import optimal_bins
+
+from .base_slicer import BaseSlicer
+
+
+class OneDSlicer(BaseSlicer):
+    """OneD Slicer allows the slicing of data into bins in a single dimension.
+
+    Parameters
+    ----------
+    slice_col_name : `str`
+        The name of the data column to base slicing on (i.e. 'airmass', etc.)
+    slice_col_units : `str`, optional
+        Set a name for the units of the sliceCol. Used for plotting labels.
+    bins : np.ndarray, optional
+        The data will be sliced into 'bins': this can be defined as an
+        array here. Default None.
+    bin_min : `float`, optional
+    bin_max : `float`, optional
+    bin_size : `float`, optional
+        If bins is not defined, then bin_min/bin_max/bin_size can be chosen
+        to anchor the slice points.
+        Default None.
+        Priority goes: bins >> bin_min/bin_max/bin_size >> data values
+        (if none of the above are chosen).
+
+    Notes
+    -----
+    All bins except for the last bin are half-open ([a, b)) while the
+    last bin is ([a, b]).
+    """
+
+    def __init__(
+        self,
+        slice_col_name=None,
+        slice_col_units=None,
+        bins=None,
+        bin_min=None,
+        bin_max=None,
+        bin_size=None,
+        verbose=True,
+        badval=0,
+    ):
+        super().__init__(verbose=verbose, badval=badval)
+        if slice_col_name is None:
+            raise ValueError("slice_col_name cannot be left None - choose a data column to group data by")
+        self.slice_col_name = slice_col_name
+        self.columns_needed = [slice_col_name]
+        # We could try to set up the self.bins here --
+        # but it's also possible that
+        # these bin_min/max/size values have not been set and
+        # should just be set from the data.
+        self.bins = bins
+        self.bin_min = bin_min
+        self.bin_max = bin_max
+        self.bin_size = bin_size
+        # Forget binmin/max/stepsize if bins was set
+        if self.bins is not None:
+            if bin_min is not None or bin_max is not None or bin_size is not None:
+                warnings.warning(
+                    f"Both bins and one of the bin_min/bin_max/bin_size was specified. "
+                    f"Using bins ({self.bins} values only."
+                )
+                self.bin_min = None
+                self.bin_max = None
+                self.bin_size = None
+        # Set the column units
+        if slice_col_units is not None:
+            self.slice_col_units = slice_col_units
+        # Try to determine the column units
+        else:
+            co = ColInfo()
+            self.slice_col_units = co.get_units(self.slice_col_name)
+        # Set slicer re-initialize values and default plotFunction
+        self.slicer_init = {
+            "slice_col_name": self.slice_col_name,
+            "slice_col_units": slice_col_units,
+            "badval": badval,
+            "bin_min": self.bin_min,
+            "bin_max": self.bin_max,
+            "bin_size": self.bin_size,
+            "bins": self.bins,
+        }
+        self.plot_funcs = [
+            OneDBinnedData,
+        ]
+
+    def setup_slicer(self, sim_data, maps=None):
+        """Set up bins in slicer.
+
+        This happens AFTER sim_data is defined,
+        thus typically in the MetricBundleGroup.
+        This maps data into the bins;
+        it's not a good idea to reuse a OneDSlicer as a result.
+        """
+        if "bins" in self.slice_points:
+            warning_msg = "Warning: this OneDSlicer was already set up once. "
+            warning_msg += (
+                "Re-setting up a OneDSlicer is unpredictable; at the very least, it "
+                "will change the mapping of the simulated data into the data slices. "
+                "A safer choice is to use a separate OneDSlicer for each MetricBundle."
+            )
+            warnings.warn(warning_msg)
+        slice_col = sim_data[self.slice_col_name]
+        # Set bins from data or specified values,
+        # if they were previously defined.
+        if self.bins is None:
+            # Set bin min/max values (could have been set in __init__)
+            if self.bin_min is None:
+                self.bin_min = np.nanmin(slice_col)
+            if self.bin_max is None:
+                self.bin_max = np.nanmax(slice_col)
+            # Give warning if bin_min = bin_max,
+            # and do something at least slightly reasonable.
+            if self.bin_min == self.bin_max:
+                warnings.warn(
+                    "bin_min = bin_max (maybe your data is single-valued?). "
+                    "Increasing bin_max by 1 (or 2*bin_size, if bin_size was set)."
+                )
+                if self.bin_size is not None:
+                    self.bin_max = self.bin_max + 2 * self.bin_size
+                else:
+                    self.bin_max = self.bin_max + 1
+            if self.bin_size is None:
+                bins = optimal_bins(slice_col, self.bin_min, self.bin_max)
+                nbins = np.round(bins)
+                self.bin_size = (self.bin_max - self.bin_min) / float(nbins)
+            # Set bins
+            self.bins = np.arange(self.bin_min, self.bin_max + self.bin_size / 2.0, self.bin_size, "float")
+        # nslice is used to stop iteration and should
+        # reflect the usable length of the bins
+        self.nslice = len(self.bins) - 1
+        # and "shape" refers to the length of the datavalues,
+        # and should be one less than # of bins because last
+        # binvalue is RH edge only
+        self.shape = self.nslice
+        # Set slice_point metadata.
+        self.slice_points["sid"] = np.arange(self.nslice)
+        self.slice_points["bins"] = self.bins
+        # Add metadata from map if needed.
+        self._run_maps(maps)
+
+        indxs = np.argsort(sim_data[self.slice_col_name])
+        data_sorted = sim_data[self.slice_col_name][indxs]
+
+        # Setting up slices such that left_edge <= data < right_edge
+        # in each slice.
+        left = np.searchsorted(data_sorted, self.bins[0:-1], "left")
+        right = np.searchsorted(data_sorted, self.bins[1:], "left")
+
+        self.sim_idxs = [indxs[le:ri] for le, ri in zip(left, right)]
+
+        # Set up _slice_sim_data method for this class.
+        @wraps(self._slice_sim_data)
+        def _slice_sim_data(islice):
+            """Slice sim_data on oneD sliceCol, to return relevant
+            indexes for slice_point.
+            """
+            idxs = self.sim_idxs[islice]
+            bin_left = self.bins[islice]
+            bin_right = self.bins[islice + 1]
+            return {
+                "idxs": idxs,
+                "slice_point": {
+                    "sid": islice,
+                    "bin_left": bin_left,
+                    "bin_right": bin_right,
+                },
+            }
+
+        setattr(self, "_slice_sim_data", _slice_sim_data)
+
+    def __eq__(self, other_slicer):
+        """Evaluate if slicers are equivalent."""
+        result = False
+        if isinstance(other_slicer, OneDSlicer):
+            if self.slice_col_name == other_slicer.slice_col_name:
+                # If slicer restored from disk or setup,
+                # then 'bins' in slice_points dict.
+                # This is preferred method to see if slicers are equal.
+                if ("bins" in self.slice_points) & ("bins" in other_slicer.slice_points):
+                    result = np.array_equal(other_slicer.slice_points["bins"], self.slice_points["bins"])
+                # However, before we 'setup' the slicer with data,
+                # the slicers could be equivalent.
+                else:
+                    if (self.bins is not None) and (other_slicer.bins is not None):
+                        result = np.array_equal(self.bins, other_slicer.bins)
+                    elif (
+                        (self.bin_size is not None)
+                        and (self.bin_min is not None) & (self.bin_max is not None)
+                        and (other_slicer.bin_size is not None)
+                        and (other_slicer.bin_min is not None)
+                        and (other_slicer.bin_max is not None)
+                    ):
+                        if (
+                            (self.bin_size == other_slicer.bin_size)
+                            and (self.bin_min == other_slicer.bin_min)
+                            and (self.bin_max == other_slicer.bin_max)
+                        ):
+                            result = True
+        return result