Merge branch 'develop'

asteca/__init__.py

@@ -1,7 +1,7 @@
-from .isochrones import isochrones
-from .synthetic import synthetic
-from .cluster import cluster
-from .likelihood import likelihood
+from .isochrones import isochrones as isochrones
+from .synthetic import synthetic as synthetic
+from .cluster import cluster as cluster
+from .likelihood import likelihood as likelihood
 
 from contextlib import suppress
 import importlib.metadata

asteca/cluster.py

@@ -12,7 +12,8 @@ class cluster:
     Parameters
     ----------
     cluster_df : pd.DataFrame
-        pandas DataFrame with the cluster's loaded data
+        `pandas DataFrame <https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html>`_
+        with the cluster's loaded data
     magnitude : str
         Name of the DataFrame column that contains the magnitude
     e_mag : str
@@ -37,6 +38,7 @@ class cluster:
         Name of the DataFrame column that contains the second color's uncertainty
 
     """
+
     cluster_df: pd.DataFrame
     magnitude: str
     e_mag: str
@@ -59,7 +61,7 @@ def _load(self):
         The photometry is store with a '_p' to differentiate from the self.magnitude,
         self.color, etc that are defined with the class is called.
         """
-        print("Reading and processing cluster data")
+        print("Instantiating cluster...")
 
         self.mag_p = np.array(self.cluster_df[self.magnitude])
         self.e_mag_p = np.array(self.cluster_df[self.e_mag])
@@ -75,6 +77,11 @@ def _load(self):
             self.ra_v = self.cluster_df[self.ra]
             self.dec_v = self.cluster_df[self.dec]
 
+        print(f"N_stars        : {len(self.mag_p)}")
+        print(f"Magnitude      : {self.magnitude}")
+        print(f"Color          : {self.color}")
+        if self.color2 is not None:
+            print(f"Color2         : {self.color2}")
         print("Cluster object generated\n")
 
     def radecplot(self):
@@ -110,57 +117,67 @@ def radecplot(self):
 
         return ax
 
-    def clustplot(self, ax=None, binar_prob=None):
+    def clustplot(self, ax, color_idx=0, binar_probs=None):
         r"""Generate a color-magnitude plot.
 
         Parameters
         ----------
         ax : matplotlib.axis, optional, default=None
             Matplotlib axis where to draw the plot
-        binar_prob : numpy.array, optional, default=None
+        color_idx : int, default=0
+            Index of the color to plot. If ``0`` (default), plot the first color. If
+            ``1`` plot the second color.
+        binar_probs : numpy.array, optional, default=None
             Array with probabilities of being a binary system for each observed star
 
         Returns
         -------
         matplotlib.axis
             Matplotlib axis object
 
-
         """
-        if ax is None:
-            f, ax = plt.subplots()
+        if color_idx > 1:
+            raise ValueError(
+                f"Wrong 'color_idx' value ({color_idx}), should be one of: [0, 1]"
+            )
 
-        if binar_prob is not None:
-            msk_binar = binar_prob > 0.5
+        if binar_probs is not None:
+            msk_binar = binar_probs > 0.5
 
         mag_col = self.magnitude
         col_col = self.color
+        if color_idx == 1:
+            col_col = self.color2
 
-        if binar_prob is None:
+        if binar_probs is None:
             ax.scatter(
-                self.colors_p[0],
+                self.colors_p[color_idx],
                 self.mag_p,
                 c="green",
                 alpha=0.5,
                 label=f"Observed, N={len(self.mag_p)}",
             )
         else:
             ax.scatter(
-                self.colors_p[0][~msk_binar],
+                self.colors_p[color_idx][~msk_binar],
                 self.mag_p[~msk_binar],
-                c="green",
+                # c="green",
+                c=binar_probs[~msk_binar],
+                marker="o",
                 alpha=0.5,
                 label=f"Observed (single), N={len(self.mag_p[~msk_binar])}",
             )
             ax.scatter(
-                self.colors_p[0][msk_binar],
+                self.colors_p[color_idx][msk_binar],
                 self.mag_p[msk_binar],
-                c="red",
+                # c="red",
+                c=binar_probs[msk_binar],
+                marker="s",
                 alpha=0.5,
                 label=f"Observed (binary), N={len(self.mag_p[msk_binar])}",
             )
 
-        ax.set_ylim(max(self.mag_p) + .5, min(self.mag_p) - .5)
+        ax.set_ylim(max(self.mag_p) + 0.5, min(self.mag_p) - 0.5)
         ax.set_xlabel(col_col)
         ax.set_ylabel(mag_col)
         ax.legend()

asteca/isochrones.py

@@ -1,4 +1,5 @@
 import os
+import numpy as np
 from dataclasses import dataclass
 from typing import Optional
 from .modules import isochrones_priv
@@ -10,39 +11,47 @@ class isochrones:
 
     This object contains the loaded theoretical isochrones used by the
     :py:mod:`asteca.synthetic` class to generate synthetic clusters.
+    See :ref:`isochronesload` for more details.
 
     Parameters
     ----------
-    isochs_path : str
-        Path to the folder that contains the files for the theoretical isochrones.
-        The name of the folder must be one of the supported isochrone services:
+    model : str, {"PARSEC", "MIST", "BASTI"}
+        The model must be one of the supported isochrone services:
         `PARSEC <http://stev.oapd.inaf.it/cgi-bin/cmd_3.7>`_,
         `MIST <https://waps.cfa.harvard.edu/MIST/>`_, or
         `BASTI <http://basti-iac.oa-abruzzo.inaf.it/isocs.html>`_.
-        Examples of valid paths: ``isochrones/PARSEC/``, ``Mist/``, ``basti``.
-        See :ref:`isochronesload` for more detailed information on how to properly
-        store the isochrone files.
-    magnitude : dict
-        Dictionary containing the magnitude's filter name (as defined in the files
-        of the theoretical isochrones) as the key, and its effective lambda
-        (in Angstrom) as the value. Example for Gaia's ``G`` magnitude:
-        ``{"Gmag": 6390.21}``.
-    color : dict
-        Dictionary containing the color used in the cluster's analysis. The correct
-        format is: ``{"filter1": 1111.11, "filter2": 2222.22}``, where ``filter1``
-        and `filter2` are the names of the filters that are combined to generate
+    isochs_path : str
+        Path to the folder that contains the files for the theoretical isochrones.
+    magnitude : str
+        Magnitude's filter name as defined in the theoretical isochrones.
+        Example for Gaia's ``G`` magnitude: ``"Gmag"``.
+    color : tuple
+        Tuple containing the filter names that generate the first color defined.
+        The correct format is: ``("filter1", "filter2")``, where ``filter1``
+        and ``filter2`` are the names of the filters that are combined to generate
         the color. The order is important because the color will be generated as:
-        ``filter1-filter2``. The values ``1111.11`` and ``2222.22`` are the effective
-        lambdas (in Angstrom) for each filter. The color does not need to be defined
-        in the same photometric system as the magnitude.
-        Example for Gaia's 'BP-RP' color:
-        ``{"G_BPmag": 5182.58, "G_RPmag": 7825.08}``
+        ``filter1-filter2``. Example for Gaia's 'BP-RP' color:
+        ``("G_BPmag", "G_RPmag")``.
+    color2 : tuple, optional, default=None
+        Optional second color to use in the analysis. Same format as that used by the
+        ``color`` parameter.
+    magnitude_effl : float, optional, default=None
+        Effective lambda (in Angstrom) for the magnitude filter.
+    color_effl : tuple, optional, default=None
+        Effective lambdas for the filters that make up the ``color`` defined in the
+        :py:mod:`asteca.isochrones` object. E.g.: ``(1111.11, 2222.22)`` where
+        ``1111.11`` and ``2222.22`` are the effective lambdas (in Angstrom) for each
+        filter, in the same order as ``color``.
+    color2_effl : tuple, optional, default=None
+        Same as ``color_effl`` but for a second (optional) color defined.
+    z_to_FeH : float, optional, default=None
+        If ``None``, the default ``z`` values in the isochrones will be used to
+        generate the synthetic clusters. If ``float``, it must represent the solar
+        metallicity for these isochrones. The metallicity values will then be converted
+        to ``[FeH]`` values, to be used by the :meth:`synthetic.generate()` method.
     N_interp : int, default=2500
         Number of interpolation points used to ensure that all isochrones are the
         same shape.
-    color2 : dict, optional, default=None
-        Optional second color to use in the analysis. Same format as that used by the
-        ``color`` parameter.
     column_names : dict, optional, default=None
         Column names for the initial mass, metallicity, and age for the photometric
         system's isochrones files. Example:
@@ -56,22 +65,33 @@ class fails to load the files.
         "`in preparation <http://stev.oapd.inaf.it/cmd_3.7/faq.html>`_".
 
     """
+
+    model: str
     isochs_path: str
-    magnitude: dict
-    color: dict
+    magnitude: str
+    color: tuple
+    color2: Optional[tuple] = None
+    magnitude_effl: Optional[float] = None
+    color_effl: Optional[tuple] = None
+    color2_effl: Optional[tuple] = None
     N_interp: int = 2500
-    color2: Optional[dict] = None
+    z_to_FeH: Optional[float] = None
     column_names: Optional[dict] = None
     parsec_rm_stage_9: Optional[bool] = True
 
     def __post_init__(self):
-        # Extract model from isochrones' path
-        if self.isochs_path.endswith("/"):
-            self.model = self.isochs_path[:-1].split("/")[-1].upper()
-        else:
-            self.model = self.isochs_path.split("/")[-1].upper()
+        # Check that the number of colors match
+        if self.color2 is not None and self.color2_effl is None:
+            raise ValueError(
+                "Second color is defined but its effective lambdas are missing."
+            )
+        if self.color2 is None and self.color2_effl is not None:
+            raise ValueError(
+                "Lambdas for the second color are defined but second color is missing."
+            )
 
         # Check model input
+        self.model = self.model.upper()
         models = ("PARSEC", "MIST", "BASTI")
         if self.model not in models:
             raise ValueError(
@@ -80,22 +100,62 @@ def __post_init__(self):
 
         # Check path to isochrones
         if os.path.isdir(self.isochs_path) is False:
-            raise ValueError(
-                f"Path '{self.isochs_path}' must point to the folder that "
-                + f"contains the '{self.model}' isochrones"
-            )
-
-        # Extract magnitude, color(s), and lambdas
-        self.mag_color_lambdas = self.magnitude | self.color
-        self.mag_filter_name = list(self.magnitude.keys())[0]
-        self.color_filter_name = [list(self.color.keys())]
-        # Add second color if any
-        if self.color2 is not None:
-            self.color_filter_name.append(list(self.color2.keys()))
-            self.mag_color_lambdas = self.mag_color_lambdas | self.color2
+            raise ValueError(f"Path '{self.isochs_path}' not found")
 
+        print("Instantiating isochrones...")
         # Load isochrone files
-        self.theor_tracks, self.color_filters, self.met_age_dict = isochrones_priv.load(
-            self
+        self.theor_tracks, self.color_filters, self.met_age_dict, N_isoch_files = (
+            isochrones_priv.load(self)
         )
+
+        # Convert z to FeH if requested
+        met_n = "z  "
+        if self.z_to_FeH is not None:
+            self._func_z_to_FeH(self.z_to_FeH)
+            met_n = "FeH"
+        zmin, zmax, amin, amax = self._min_max()
+
+        N_met, N_age, N_cols, N_interp = self.theor_tracks.shape
+        print(f"Model          : {self.model}")
+        print(f"N_files        : {N_isoch_files}")
+        print(f"N_met          : {N_met}")
+        print(f"N_age          : {N_age}")
+        print(f"N_interp       : {N_interp}")
+        print(f"{met_n} range      : [{zmin}, {zmax}]")
+        print(f"loga range     : [{amin}, {amax}]")
+        print(f"Magnitude      : {self.magnitude}")
+        print(f"Color          : {self.color[0]}-{self.color[1]}")
+        if self.color2 is not None:
+            print(f"Color2         : {self.color2[0]}-{self.color2[1]}")
         print("Isochrone object generated\n")
+
+    def _func_z_to_FeH(self, z_to_FeH):
+        r"""Convert z to FeH"""
+        feh = np.log10(self.met_age_dict["met"] / z_to_FeH)
+        N_old = len(feh)
+        round_n = 4
+        while True:
+            feh_r = np.round(feh, round_n)
+            N_new = len(set(feh_r))
+            # If no duplicated values exist after rounding
+            if N_old == N_new:
+                break
+            round_n += 1
+        # Replace old values
+        self.met_age_dict["met"] = feh_r
+
+    def _min_max(self) -> tuple[float]:
+        r"""Return the minimum and maximum values for the metallicity and age defined
+        in the theoretical isochrones.
+
+        Returns
+        -------
+        tuple[float]
+            Tuple of (minimum_metallicity, maximum_metallicity, minimum_age, maximum_age)
+
+        """
+        zmin = self.met_age_dict["met"].min()
+        zmax = self.met_age_dict["met"].max()
+        amin = self.met_age_dict["loga"].min()
+        amax = self.met_age_dict["loga"].max()
+        return zmin, zmax, amin, amax

asteca/likelihood.py

@@ -28,6 +28,7 @@ class likelihood:
         the color(s), also with edge values.
 
     """
+
     my_cluster: cluster
     lkl_name: str = "plr"
     bin_method: str = "knuth"

asteca/modules/imfs.py

@@ -39,7 +39,7 @@ def sampled_inv_cdf(ijkseed, N):
         # mr = np.random.rand(N)
         # The internal seed can not be 'self.seed' because otherwise all the
         # floats returned are equal
-        mr = np.random.default_rng(ijkseed).uniform(0., 1., N)
+        mr = np.random.default_rng(ijkseed).uniform(0.0, 1.0, N)
         return inv_cdf(mr)
 
     # Sample in chunks until the maximum defined mass is reached.