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

rubin_sim/maf/db/tracking_db.py

@@ -0,0 +1,332 @@
+__all__ = ("TrackingDb", "add_run_to_database")
+
+import os
+
+from sqlalchemy import Column, Integer, String, create_engine
+from sqlalchemy.engine import url
+from sqlalchemy.exc import DatabaseError
+from sqlalchemy.orm import declarative_base, sessionmaker
+
+from . import ResultsDb, VersionRow
+
+Base = declarative_base()
+
+
+class RunRow(Base):
+    """
+    Define contents and format of run list table.
+
+    Table to list all available MAF results,
+    along with their opsim run and some comment info.
+    """
+
+    __tablename__ = "runs"
+    # Define columns in metric list table.
+    maf_run_id = Column(Integer, primary_key=True)
+    run_group = Column(String)
+    run_name = Column(String)
+    run_comment = Column(String)
+    run_version = Column(String)
+    run_date = Column(String)
+    db_file = Column(String)
+    maf_comment = Column(String)
+    maf_version = Column(String)
+    maf_date = Column(String)
+    maf_dir = Column(String)
+
+    def __repr__(self):
+        rstr = (
+            "<Run(maf_run_id='%d', run_group='%s', run_name='%s', run_comment='%s', "
+            "run_version='%s', run_date='%s', maf_comment='%s', "
+            "maf_version='%s', maf_date='%s', maf_dir='%s', db_file='%s'>"
+            % (
+                self.maf_run_id,
+                self.run_group,
+                self.run_name,
+                self.maf_comment,
+                self.maf_version,
+                self.run_date,
+                self.maf_comment,
+                self.maf_version,
+                self.maf_date,
+                self.maf_dir,
+                self.db_file,
+            )
+        )
+        return rstr
+
+
+class TrackingDb:
+    """Sqlite database to track MAF output runs and their locations,
+    for show_maf
+    """
+
+    def __init__(self, database=None, trackingDbverbose=False):
+        """
+        Set up the Tracking Database.
+        """
+        self.verbose = trackingDbverbose
+        self.driver = "sqlite"
+        # connecting to non-existent database creates it automatically
+        if database is None:
+            # Default is a file in the current directory.
+            self.database = "trackingDb_sqlite.db"
+            self.tracking_db_dir = "."
+        else:
+            self.database = database
+            self.tracking_db_dir = os.path.dirname(database)
+        # only sqlite
+        dbAddress = url.URL.create(drivername=self.driver, database=self.database)
+        engine = create_engine(dbAddress, echo=self.verbose)
+        if self.verbose:
+            print("Created or connected to MAF tracking %s database at %s" % (self.driver, self.database))
+        self.Session = sessionmaker(bind=engine)
+        self.session = self.Session()
+        # Create the tables, if they don't already exist.
+        try:
+            Base.metadata.create_all(engine)
+        except DatabaseError:
+            raise DatabaseError(
+                "Cannot create a %s database at %s. Check directory exists." % (self.driver, self.database)
+            )
+
+    def close(self):
+        self.session.close()
+
+    def add_run(
+        self,
+        run_group=None,
+        run_name=None,
+        run_comment=None,
+        run_version=None,
+        run_date=None,
+        maf_comment=None,
+        maf_version=None,
+        maf_date=None,
+        maf_dir=None,
+        db_file=None,
+        maf_run_id=None,
+    ):
+        """Add a run to the tracking database.
+
+        Parameters
+        ----------
+        run_group : `str`, optional
+            Set a name to group this run with (eg. "Tier 1, 2016").
+        run_name : `str`, optional
+            Set a name for the opsim run.
+        run_comment : `str`, optional
+            Set a comment describing the opsim run.
+        run_version : `str`, optional
+            Set the version of opsim.
+        run_date : `str`, optional
+            Set the date the opsim run was created.
+        maf_comment : `str`, optional
+            Set a comment to describe the MAF analysis.
+        maf_version : `str`, optional
+            Set the version of MAF used for analysis.
+        maf_date : `str`, optional
+            Set the date the MAF analysis was run.
+        maf_dir : `str`, optional
+            The relative path to the MAF directory.
+            Will be converted to a relative path if absolute.
+        db_file : `str`, optional
+            The relative path to the Opsim SQLite database file.
+        maf_run_id : `int`, optional
+            The maf_run_id to assign to this record in the database
+            (note this is a primary key!).
+            If this run (ie the maf_dir) exists in the database already,
+            this will be ignored.
+
+        Returns
+        -------
+        maf_run_id : `int`
+            The maf_run_id stored in the database.
+        """
+        if run_group is None:
+            run_group = "NULL"
+        if run_name is None:
+            run_name = "NULL"
+        if run_comment is None:
+            run_comment = "NULL"
+        if run_version is None:
+            run_version = "NULL"
+        if run_date is None:
+            run_date = "NULL"
+        if maf_comment is None:
+            maf_comment = "NULL"
+        if maf_version is None:
+            maf_version = "NULL"
+        if maf_date is None:
+            maf_date = "NULL"
+        if maf_dir is None:
+            maf_dir = "NULL"
+        if maf_dir is not None:
+            maf_dir = os.path.relpath(maf_dir, start=self.tracking_db_dir)
+        if db_file is None:
+            db_file = "NULL"
+        # Test if maf_dir already exists in database.
+        prevrun = self.session.query(RunRow).filter_by(maf_dir=maf_dir).all()
+        if len(prevrun) > 0:
+            runIds = []
+            for run in prevrun:
+                runIds.append(run.maf_run_id)
+            print(
+                "Updating information in tracking database - %s already present with runId %s."
+                % (maf_dir, runIds)
+            )
+            for run in prevrun:
+                self.session.delete(run)
+            self.session.commit()
+            runinfo = RunRow(
+                maf_run_id=runIds[0],
+                run_group=run_group,
+                run_name=run_name,
+                run_comment=run_comment,
+                run_version=run_version,
+                run_date=run_date,
+                maf_comment=maf_comment,
+                maf_version=maf_version,
+                maf_date=maf_date,
+                maf_dir=maf_dir,
+                db_file=db_file,
+            )
+        else:
+            if maf_run_id is not None:
+                # Check if maf_run_id exists already.
+                existing = self.session.query(RunRow).filter_by(maf_run_id=maf_run_id).all()
+                if len(existing) > 0:
+                    raise ValueError(
+                        "maf_run_id %d already exists in database, for %s. "
+                        "Record must be deleted first." % (maf_run_id, existing[0].maf_dir)
+                    )
+                runinfo = RunRow(
+                    maf_run_id=maf_run_id,
+                    run_group=run_group,
+                    run_name=run_name,
+                    run_comment=run_comment,
+                    run_version=run_version,
+                    run_date=run_date,
+                    maf_comment=maf_comment,
+                    maf_version=maf_version,
+                    maf_date=maf_date,
+                    maf_dir=maf_dir,
+                    db_file=db_file,
+                )
+            else:
+                runinfo = RunRow(
+                    run_group=run_group,
+                    run_name=run_name,
+                    run_comment=run_comment,
+                    run_version=run_version,
+                    run_date=run_date,
+                    maf_comment=maf_comment,
+                    maf_version=maf_version,
+                    maf_date=maf_date,
+                    maf_dir=maf_dir,
+                    db_file=db_file,
+                )
+        self.session.add(runinfo)
+        self.session.commit()
+        return runinfo.maf_run_id
+
+    def delRun(self, runId):
+        """
+        Remove a run from the tracking database.
+        """
+        runinfo = self.session.query(RunRow).filter_by(maf_run_id=runId).all()
+        if len(runinfo) == 0:
+            raise Exception("Could not find run with maf_run_id %d" % (runId))
+        if len(runinfo) > 1:
+            raise Exception("Found more than one run with maf_run_id %d" % (runId))
+        print("Removing run info for runId %d " % (runId))
+        print(" ", runinfo)
+        self.session.delete(runinfo[0])
+        self.session.commit()
+
+
+def add_run_to_database(
+    maf_dir,
+    tracking_db_file,
+    run_group=None,
+    run_name=None,
+    run_comment=None,
+    maf_comment=None,
+    db_file=None,
+    maf_version=None,
+    maf_date=None,
+    sched_version=None,
+    sched_date=None,
+    skip_extras=False,
+):
+    """Adds information about a MAF analysis run to a MAF tracking database.
+
+    Parameters
+    ----------
+    maf_dir : `str`
+        Path to the directory where the MAF results are located.
+    tracking_db_file : `str`
+        Full filename (+path) to the tracking database to use.
+    run_group: `str`, optional
+        Name to use to group this run with other opsim runs.
+    run_name : `str`, optional
+        Name of the opsim run.
+    run_comment : `str`, optional
+        Comment about the opsim run.
+    run_version : `str`, optional
+        Value to use for the opsim version information.
+    maf_comment : `str`, optional
+        Comment about the MAF analysis.
+    db_file : `str`, optional
+        Relative path + name of the opsim database file.
+    """
+    trackingDb = TrackingDb(database=tracking_db_file)
+
+    maf_dir = os.path.relpath(maf_dir, start=os.path.dirname(trackingDb.tracking_db_dir))
+    if not os.path.isdir(maf_dir):
+        raise ValueError("There is no directory containing MAF outputs at %s." % (maf_dir))
+
+    # Connect to resultsDb for additional information if available
+    if os.path.isfile(os.path.join(maf_dir, "resultsDb_sqlite.db")) and not skip_extras:
+        resdb = ResultsDb(maf_dir)
+        if run_name is None:
+            run_name = resdb.get_run_name()
+            if len(run_name) > 1:
+                run_name = 0
+            else:
+                run_name = run_name[0]
+        if maf_version is None or maf_date is None:
+            resdb.open()
+            query = resdb.session.query(VersionRow).all()
+            for v in query:
+                if maf_version is None:
+                    maf_version = v.version
+                if maf_date is None:
+                    maf_date = v.run_date
+            resdb.close()
+
+    print("Adding to tracking database at %s:" % (tracking_db_file))
+    print(" Maf_dir = %s" % (maf_dir))
+    print(" Maf_comment = %s" % (maf_comment))
+    print(" run_group = %s" % (run_group))
+    print(" run_name = %s" % (run_name))
+    print(" run_comment = %s" % (run_comment))
+    print(" run_version = %s" % (sched_version))
+    print(" run_date = %s" % (sched_date))
+    print(" maf_version = %s" % (maf_version))
+    print(" maf_date = %s" % (maf_date))
+    print(" db_file = %s" % (db_file))
+    runId = trackingDb.add_run(
+        run_group=run_group,
+        run_name=run_name,
+        run_comment=run_comment,
+        run_version=sched_version,
+        run_date=sched_date,
+        maf_comment=maf_comment,
+        maf_version=maf_version,
+        maf_date=maf_date,
+        maf_dir=maf_dir,
+        db_file=db_file,
+    )
+    print("Used MAF RunID %d" % (runId))
+    trackingDb.close()

rubin_sim/maf/ddf_dir.py

@@ -0,0 +1,53 @@
+import argparse
+import glob
+import os
+import shutil
+
+import rubin_sim.maf.batches as batches
+import rubin_sim.maf.db as db
+import rubin_sim.maf.metric_bundles as mb
+
+__all__ = ("ddf_dir",)
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+
+def ddf_dir():
+    """
+    Run the glance batch on all .db files in a directory.
+    """
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--db", type=str, default=None)
+    parser.add_argument("--nside", type=int, default=512)
+    parser.add_argument("--old_coords", dest="old_coords", action="store_true")
+    parser.set_defaults(verbose=False)
+
+    args = parser.parse_args()
+
+    if args.db is None:
+        db_files = glob.glob("*.db")
+        db_files = [filename for filename in db_files if "trackingDb" not in filename]
+    else:
+        db_files = [args.db]
+    run_names = [os.path.basename(name).replace(".db", "") for name in db_files]
+
+    for filename, name in zip(db_files, run_names):
+        if os.path.isdir(name + "_ddf"):
+            shutil.rmtree(name + "_ddf")
+
+        bdict = {}
+        bdict.update(batches.ddfBatch(run_name=name, nside=args.nside, old_coords=args.old_coords))
+        results_db = db.ResultsDb(out_dir=name + "_ddf")
+        group = mb.MetricBundleGroup(
+            bdict,
+            filename,
+            out_dir=name + "_ddf",
+            results_db=results_db,
+            save_early=False,
+        )
+        group.run_all(clear_memory=True, plot_now=True)
+        results_db.close()
+        db.add_run_to_database(name + "_ddf", "trackingDb_sqlite.db", None, name, "", "", name + ".db")

rubin_sim/maf/generate_ss.py

@@ -0,0 +1,218 @@
+__all__ = ("generate_ss_commands", "generate_ss")
+
+import argparse
+import glob
+import os
+import sqlite3
+
+import numpy as np
+import pandas as pd
+
+import rubin_sim.maf.db as db
+from rubin_sim.data import get_data_dir
+
+
+def generate_ss_commands(
+    dbfiles=None,
+    pops=None,
+    split=False,
+):
+    if dbfiles is None:
+        dbfiles = glob.glob("*.db")
+        dbfiles.sort()
+
+    data_dir = os.path.join(get_data_dir(), "orbits")
+    orbit_files = {
+        "granvik_5k": os.path.join(data_dir, "granvik_5k.txt"),
+        "granvik_pha_5k": os.path.join(data_dir, "granvik_pha_5k.txt"),
+        "vatiras_granvik_10k": os.path.join(data_dir, "vatiras_granvik_10k.txt"),
+        "l7_5k": os.path.join(data_dir, "l7_5k.txt"),
+        "mba_5k": os.path.join(data_dir, "mba_5k.txt"),
+        "occ_rmax5_5k": os.path.join(data_dir, "occ_rmax5_5k.txt"),
+        "occ_rmax20_5k": os.path.join(data_dir, "occ_rmax20_5k.txt"),
+        "trojan_5k": os.path.join(data_dir, "trojan_5k.txt"),
+    }
+
+    pos_dir = os.path.join(get_data_dir(), "orbits_precompute")
+    position_files = {
+        "granvik_5k": os.path.join(pos_dir, "granvik_5k.npz"),
+        "granvik_pha_5k": os.path.join(pos_dir, "granvik_pha_5k.npz"),
+        "vatiras_granvik_10k": os.path.join(pos_dir, "vatiras_granvik_10k.npz"),
+        "l7_5k": os.path.join(pos_dir, "l7_5k.npz"),
+        "mba_5k": os.path.join(pos_dir, "mba_5k.npz"),
+        "occ_rmax5_5k": os.path.join(pos_dir, "occ_rmax5_5k.npz"),
+        "occ_rmax20_5k": os.path.join(pos_dir, "occ_rmax20_5k.npz"),
+        "trojan_5k": os.path.join(pos_dir, "trojan_5k.npz"),
+    }
+
+    objtypes = {
+        "granvik_5k": "NEO",
+        "granvik_pha_5k": "PHA",
+        "vatiras_granvik_10k": "Vatira",
+        "l7_5k": "TNO",
+        "mba_5k": "MBA",
+        "occ_rmax5_5k": "OCC_r5",
+        "occ_rmax20_5k": "OCC_r20",
+        "sdo_5k": "SDO",
+        "trojan_5k": "Trojan",
+    }
+
+    if pops is None:
+        # put in order so longer runtime ones are first
+        pops = [
+            "trojan_5k",
+            "l7_5k",
+            "mba_5k",
+            "granvik_5k",
+            "granvik_pha_5k",
+            "vatiras_granvik_10k",
+            "occ_rmax5_5k",
+            "occ_rmax20_5k",
+        ]
+
+    elif pops is not None:
+        pp = [p for p in orbit_files.keys() if p == pops]
+        if len(pp) == 0:
+            print(f"Did not find population {pops} in expected types ({list(orbit_files.keys())}")
+        pops = [pops]
+
+    runs = [os.path.split(file)[-1].replace(".db", "") for file in dbfiles]
+    runs = [run for run in runs if "tracking" not in run]
+
+    # Find mjd_start for each run:
+    mjd_starts = {}
+    for filename in dbfiles:
+        con = sqlite3.connect(filename)
+        value = pd.read_sql("select min(observationStartMJD) from observations;", con)
+        con.close()
+        mjd_starts[filename] = np.floor(value.values.min())
+
+    if not split:
+        output_file = open("ss_script.sh", "w")
+        for run, filename in zip(runs, dbfiles):
+            out_dir = f"{run}_ss"
+            try:
+                os.mkdir(out_dir)
+            except FileExistsError:
+                pass
+            # Create the results DB so multiple threads
+            # don't try to create it later
+            # This isn't used in this script, but needs to exist on-disk.
+            results_db = db.ResultsDb(out_dir=out_dir)
+        for pop in pops:
+            for run, filename in zip(runs, dbfiles):
+                objtype = objtypes[pop]
+
+                s1 = (
+                    f"make_lsst_obs --simulation_db {filename} --orbit_file {orbit_files[pop]}"
+                    f" --positions_file {position_files[pop]}"
+                )
+                s2 = (
+                    f"run_moving_calc --obs_file {run}__{pop}_obs.npz"
+                    f" --simulation_db {filename} --orbit_file {orbit_files[pop]}"
+                    f" --out_dir {run}_ss"
+                    f" --objtype {objtype}"
+                    f" --start_time {mjd_starts[filename]}"
+                )
+                s3 = (
+                    f"run_moving_fractions --work_dir {run}_ss"
+                    f" --metadata {objtype}"
+                    f" --start_time {mjd_starts[filename]}"
+                )
+                print(s1 + " ; " + s2 + " ; " + s3, file=output_file)
+    else:
+        for run, filename in zip(runs, dbfiles):
+            out_dir = f"{run}_ss"
+            try:
+                os.mkdir(out_dir)
+            except FileExistsError:
+                pass
+            # Create the results DB so multiple threads
+            # don't try to create it later
+            results_db = db.ResultsDb(out_dir=out_dir)  # noqa F841
+            outfile = f"{run}_ss_script.sh"
+            if split:
+                output_file = open(outfile, "w")
+            for pop in pops:
+                objtype = objtypes[pop]
+                if split:
+                    splitfiles = glob.glob(os.path.join(data_dir, "split") + f"/*{pop}*")
+                    outfile_split = outfile.replace(".sh", f"_{pop}_split.sh")
+                    # If the output split file already exists, remove it
+                    # (as we append, not write)
+                    if os.path.isfile(outfile_split):
+                        os.remove(outfile_split)
+                    for i, splitfile in enumerate(splitfiles):
+                        split = os.path.split(splitfile)[-1]
+                        split = split.replace(".des", "").replace(".s3m", "").replace(".txt", "")
+                        with open(outfile_split, "a") as wi:
+                            s1 = (
+                                f"make_lsst_obs --simulation_db {filename} --orbit_file {splitfile}"
+                                f" --positions_file {position_files[pop]}"
+                                f" --out_dir {out_dir}"
+                            )
+                            s2 = (
+                                f"run_moving_calc --obs_file {out_dir}/{run}__{split}_obs.npz"
+                                f" --simulation_db {filename} --orbit_file {orbit_files[pop]}"
+                                f" --out_dir {out_dir}/{split}"
+                                f" --run_name {run}"
+                                f" --objtype {objtype}"
+                                f" --start_time {mjd_starts[filename]}"
+                            )
+                            print(s1 + " ; " + s2, file=wi)
+                    s3 = (
+                        f"run_moving_join --orbit_file {pop}"
+                        f" --base_dir {out_dir}"
+                        f" --out_dir {out_dir}/sso"
+                    )
+                    s4 = (
+                        f"run_moving_fractions --work_dir {out_dir}/sso"
+                        f" --metadata {objtype}"
+                        f" --start_time {mjd_starts[filename]}"
+                    )
+                    print(
+                        f"cat {outfile_split} | parallel -j 10 ; {s3}  ; {s4}",
+                        file=output_file,
+                    )
+    output_file.close()
+
+
+def generate_ss():
+    """Generate solar system processing commands."""
+
+    parser = argparse.ArgumentParser(description="Generate solar system processing commands")
+    parser.add_argument("--db", type=str, default=None, help="database to process")
+    parser.set_defaults(vatiras=False)
+    parser.add_argument("--pop", type=str, default=None, help="identify one population to run")
+    parser.add_argument("--start_mjd", type=float, default=None, help="start of the sim")
+    parser.add_argument(
+        "--split",
+        dest="split",
+        default=False,
+        action="store_true",
+        help="Split up population files; rejoin during processing",
+    )
+    args = parser.parse_args()
+
+    if args.db is None:
+        # Just look for any .db files in this directory
+        db_files = glob.glob("*.db")
+        # But remove trackingDb and results_db if they're there
+        try:
+            db_files.remove("trackingDb_sqlite.db")
+        except ValueError:
+            pass
+        try:
+            db_files.remove("resultsDb_sqlite.db")
+        except ValueError:
+            pass
+    elif isinstance(args.db, str):
+        db_files = [args.db]
+    else:
+        db_files = args.db
+
+    generate_ss_commands(
+        split=args.split,
+        dbfiles=db_files,
+        pops=args.pop,
+    )

rubin_sim/maf/glance_dir.py

@@ -0,0 +1,58 @@
+import argparse
+import glob
+import os
+import shutil
+
+import rubin_sim.maf.batches as batches
+import rubin_sim.maf.db as db
+import rubin_sim.maf.metric_bundles as mb
+
+__all__ = ("glance_dir",)
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+
+def glance_dir():
+    """
+    Run the glance batch on all .db files in a directory.
+    """
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--db", type=str, default=None)
+    args = parser.parse_args()
+
+    if args.db is None:
+        db_files = glob.glob("*.db")
+        db_files = [filename for filename in db_files if "trackingDb" not in filename]
+    else:
+        db_files = [args.db]
+    run_names = [os.path.basename(name).replace(".db", "") for name in db_files]
+
+    for filename, name in zip(db_files, run_names):
+        if os.path.isdir(name + "_glance"):
+            shutil.rmtree(name + "_glance")
+        colmap = batches.col_map_dict()
+
+        bdict = {}
+        bdict.update(batches.glanceBatch(colmap, name))
+        results_db = db.ResultsDb(out_dir=name + "_glance")
+        group = mb.MetricBundleGroup(
+            bdict,
+            filename,
+            out_dir=name + "_glance",
+            results_db=results_db,
+            save_early=False,
+        )
+        group.run_all(clear_memory=True, plot_now=True)
+        results_db.close()
+        db.add_run_to_database(
+            maf_dir=name + "_glance",
+            tracking_db_file="trackingDb_sqlite.db",
+            run_group=None,
+            run_name=name,
+            run_comment="",
+            maf_comment="Glance",
+            db_file=name + ".db",
+        )

rubin_sim/maf/maf_contrib/__init__.py

@@ -0,0 +1,25 @@
+# Add similar lines (from .filename import *) when you add new metrics,
+# stackers or slicers.
+from .calculate_lsst_field_visibility_astropy import *
+from .depth_limited_num_gal_metric import *
+from .filter_pair_t_gaps_metric import *
+from .grb_transient_metric import *
+from .gw170817_det_metric import *
+from .kne_metrics import *
+from .lss_metrics import *
+from .lss_obs_strategy import *
+from .lv_dwarfs import *
+from .microlensing_metric import *
+from .periodic_metric import *
+from .periodic_star_metric import *
+from .periodic_star_modulation_metric import *
+from .presto_color_kne_pop_metric import *
+from .selfcal_uniformity_metric import *
+from .star_count_mass_metric import *
+from .star_count_metric import *
+from .static_probes_fom_summary_metric import *
+from .tdes_pop_metric import *
+from .triplet_metric import *
+from .var_depth_metric import *
+from .xrb_metrics import *
+from .young_stellar_objects_metric import *