# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""
fiberassign.utils
=======================
Utility functions.
"""
from __future__ import absolute_import, division, print_function
# stdlib
from datetime import datetime
from importlib import resources
import os
import subprocess
from time import time, sleep
# external
from astropy.table import Table
from astropy.time import Time
import fitsio
import numpy as np
import yaml
# DESI etc.
from desiutil.log import get_logger
# DG - The linter will claimi the following are unused, but they are here so you can
# import them later through `from fiberassign.utils import Logger`
from ._internal import (Logger, Timer, GlobalTimers, Circle, Segments, Shape,
Environment)
log = get_logger()
# Multiprocessing environment setup
default_mp_proc = None
"""Default number of multiprocessing processes.
Set globally on first import.
"""
if "SLURM_CPUS_PER_TASK" in os.environ:
default_mp_proc = int(os.environ["SLURM_CPUS_PER_TASK"])
else:
import multiprocessing as _mp
default_mp_proc = max(1, _mp.cpu_count() // 2)
[docs]
def option_list(opts):
"""Convert key, value pairs into a list.
This converts a dictionary into an options list that can be passed to
ArgumentParser.parse_args(). The value for each dictionary key will be
converted to a string. Values that are True will be assumed to not have
a string argument added to the options list.
Args:
opts (dict): Dictionary of options.
Returns:
(list): The list of options.
"""
optlist = []
for key, val in opts.items():
keystr = "--{}".format(key)
if val is not None:
if isinstance(val, bool):
if val:
optlist.append(keystr)
else:
optlist.append(keystr)
if isinstance(val, float):
optlist.append("{:.14e}".format(val))
elif isinstance(val, (list, tuple)):
optlist.extend(val)
else:
optlist.append("{}".format(val))
return optlist
def assert_isoformat_utc(time_str):
"""
Asserts if a date formats as "YYYY-MM-DDThh:mm:ss+00:00".
Args:
time_str: string with a date
Returns:
boolean asserting if time_str formats as "YYYY-MM-DDThh:mm:ss+00:00"
"""
try:
test_time = datetime.strptime(time_str, "%Y-%m-%dT%H:%M:%S%z")
except ValueError:
return False
# AR/SB it parses as an ISO string, now just check UTC timezone +00:00 and not +0000
return time_str.endswith("+00:00")
def get_mjd(utc_time_mjd):
"""
Get the MJD value for an input date.
Args:
utc_time_mjd: date in the "YYYY-MM-DDThh:mm:ss+00:00" format (string)
Returns:
MJD value (float)
"""
assert_isoformat_utc(utc_time_mjd)
return Time(datetime.strptime(utc_time_mjd, "%Y-%m-%dT%H:%M:%S%z")).mjd
def get_date_cutoff(datetype, cutoff_case):
"""
Returns the cutoff date for a particular case, from reading the data/cutoff-dates.yaml file.
Args:
datetype: one of the key of data/cutoff-dates.yaml (e.g., "rundate", "mtltime", "pmtime_utc_str") (string)
cutoff_case: one of the config[datetype] keys in data/cutoff-dates.yaml, (e.g. "std_wd" for datetype="rundate") (string)
Returns:
date_cutoff: cutoff date for datetype and cutoff_case, in the "YYYY-MM-DDThh:mm:ss+00:00" formatting (string)
"""
fn = str(resources.files("fiberassign").joinpath("data", "cutoff-dates.yaml"))
with open(fn, "r") as f:
config = yaml.safe_load(f)
f.close()
date_cutoff = config[datetype][cutoff_case]
return date_cutoff
def get_default_static_obsdate():
"""
Returns the 'historical' default obsdate value.
Args:
None
Returns:
"2022-07-01"
"""
return "2022-07-01"
def get_obsdate(rundate=None):
"""
Returns the default obsdate: "2022-07-01" if rundate=None or rundate<obsdate_cutoff.
Args:
rundate (optional, defaults to None): rundate, in the "YYYY-MM-DDThh:mm:ss+00:00" format (string)
Returns:
obsdate: "YYYY-MM-DD" format (string)
is_after_cutoff: is rundate after rundate_cutoff? (bool)
"""
obsdate = get_default_static_obsdate()
is_after_cutoff = False
if rundate is None:
log.info(
"rundate={} -> (obsdate, is_after_cutoff)=({}, {})".format(
rundate, obsdate, is_after_cutoff
)
)
else:
if not assert_isoformat_utc(rundate):
msg = "rundate={} is not yyyy-mm-ddThh:mm:ss+00:00".format(rundate)
log.info(msg)
raise ValueError(msg)
rundate_mjd = Time(datetime.strptime(rundate, "%Y-%m-%dT%H:%M:%S%z")).mjd
rundate_cutoff = get_date_cutoff("rundate", "obsdate")
rundate_mjd_cutoff = Time(datetime.strptime(rundate_cutoff, "%Y-%m-%dT%H:%M:%S%z")).mjd
if rundate_mjd >= rundate_mjd_cutoff:
is_after_cutoff = True
yyyy = int(rundate[:4])
obsdate = "{}{}".format(yyyy + 1, rundate[4:10])
log.info(
"rundate={} >= rundate_cutoff={} -> (obsdate, is_after_cutoff)=({}, {})".format(
rundate, rundate_cutoff, obsdate, is_after_cutoff)
)
else:
log.info(
"rundate={} < rundate_cutoff={} -> (obsdate, is_after_cutoff)=({}, {})".format(
rundate, rundate_cutoff, obsdate, is_after_cutoff)
)
return obsdate, is_after_cutoff
def get_fba_use_fabs(rundate):
"""
Return the value to which set the $FIBERASSIGN_ALGORITHM_EPOCH environment variable,
which drives some way the cpp computation is done.
Args:
rundate: rundate, in the "YYYY-MM-DDThh:mm:ss+00:00" formatting (string)
Returns:
fba_use_fabs: an integer (int)
Notes:
See this PR https://github.com/desihub/fiberassign/pull/470.
In the fiberassign running function, then one will set:
os.environ["FIBERASSIGN_USE_FABS"] = fba_use_fabs
So far:
- fba_use_fabs=0 means we reproduce the buggy behavior;
- fba_use_fabs=1 means we use the expected behavior.
"""
# AR get the cutoff dates, and corresponding values
mydict = get_date_cutoff("rundate", "fba_use_fabs")
cutoff_rundates = np.array([key for key in mydict])
cutoff_mjds = np.array([get_mjd(_) for _ in cutoff_rundates])
values = np.array([mydict[_] for _ in cutoff_rundates])
log.info(
"fba_use_fabs cutoff dates: {}".format(
", ".join(["{} = {}".format(rundate, value) for rundate, value in zip(cutoff_rundates, values)])
)
)
# AR safe, order by increasing mjd
ii = cutoff_mjds.argsort()
cutoff_rundates, cutoff_mjds, values = cutoff_rundates[ii], cutoff_mjds[ii], values[ii]
# AR now pick the earliest date after the input rundate
# AR as we have set the first cutoff date as 2019-09-16,
# AR there will always have some index returned here
# AR though still checking if someone queries with an earlier rundate...
input_mjd = get_mjd(rundate)
if input_mjd < cutoff_mjds[0]:
msg = "rundate={} is earlier than DESI! ({})".format(
rundate, cutoff_rundates[0]
)
log.error(msg)
raise ValueError(msg)
i = np.where(cutoff_mjds <= get_mjd(rundate))[0][-1]
log.info("pick fba_use_fabs = {} for rundate = {}".format(values[i], rundate))
return values[i]
def get_svn_version(svn_dir):
"""
Gets the SVN revision number of an SVN folder.
Args:
svn_dir: SVN folder path (string)
Returns:
svnver: SVN revision number of svn_dir, or "unknown" if not an svn checkout
Notes:
`svn_dir` can contain environment variables to expand, e.g. "$DESIMODEL/data"
"""
cmd = ["svn", "info", os.path.expandvars(svn_dir)]
try:
svn_ver = (
subprocess.check_output(cmd, stderr=subprocess.DEVNULL).strip().decode()
)
# search for "Last Changed Rev: " line and parse out revision number. Recent versions
# of svn have a --show-item argument that does this in a less fragile way,
# but the svn installed at KPNO is old and doesn't support this option.
searchstr = 'Last Changed Rev: '
svn_ver = [line[len(searchstr):] for line in svn_ver.split('\n')
if searchstr in line][0]
except subprocess.CalledProcessError:
svn_ver = "unknown"
return svn_ver
def get_last_line(fn):
"""
Return the last line of a text file.
Args:
fn: file name (string)
Returns:
last_line: (string)
Notes:
Fails if fn has one line only; we do not protect for that case,
as this function is intended to be used in get_program_latest_timestamp()
to read *ecsv ledgers, which will always have more than one line,
and we want the fastest function possible, to use in fiberassign on-the-fly.
Copied from https://stackoverflow.com/questions/46258499/how-to-read-the-last-line-of-a-file-in-python.
"""
with open(fn, "rb") as f:
f.seek(-2, os.SEEK_END)
while f.read(1) != b"\n":
f.seek(-2, os.SEEK_CUR)
last_line = f.readline().decode().strip()
f.close()
return last_line
def read_ecsv_keys(fn):
"""
Returns the column content of an .ecsv file.
Args:
fn: filename with an .ecsv format (string)
Returns:
keys: list of the column names in fn (list)
Notes:
Gets the column names from the first line not starting with "#".
"""
keys = []
with open(fn) as f:
for line in f:
if line[0] == "#":
continue
if len(line.strip()) == 0:
continue
keys = line.split()
break
f.close()
return keys
def get_fa_gitversion():
"""
Returns `git describe --tags --dirty --always`, or "unknown" if not a git repo
Notes:
Copied (and slightly edited) from desitarget.io.gitversion()
"""
#
origdir = os.getcwd()
os.chdir(os.path.dirname(__file__))
try:
p = subprocess.Popen(["git", "describe", "--tags", "--dirty", "--always"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
except EnvironmentError:
return "unknown"
#
os.chdir(origdir)
out = p.communicate()[0]
if p.returncode == 0:
# - avoid py3 bytes and py3 unicode; get native str in both cases
return str(out.rstrip().decode("ascii"))
else:
return "unknown"
def get_revs_dates(svndir, subdirs):
"""
Returns the list of svn revisions (and dates) for a list of subfolders.
Args:
svndir: full path to a svn checkout (e.g. "$DESI_TARGET/fiberassign/tiles/trunk") (str)
subdirs: comma-separated list of subfolders in svndir (e.g. "000,001,080") (str)
Returns:
revs: list of revisions (np array of int)
dates: list of the dates of the revisions (YYYY-MM-DD) (np array of str)
Notes:
This function queries the server ("svn log"), so do not run in parallel!
"""
revs, dates = [], []
for subdir in subdirs.split(","):
start = time()
ps = subprocess.Popen(
["svn", "log", os.path.join(svndir, subdir)], stdout=subprocess.PIPE
)
output = subprocess.check_output(
["grep", "^r[1-9]"],
stdin=ps.stdout,
)
output = output.strip().decode()
ps.stdout.close()
lines = output.split("\n")
revs += [int(line.split("|")[0][1:]) for line in lines]
dates += [line.split("|")[2].split()[0] for line in lines]
dt = time() - start
log.info("subdir={}\tnrev={}\t(dt={:.1f}s)".format(subdir, len(lines), dt))
# AR protect from overloading the server
sleep(0.5)
revs, ii = np.unique(revs, return_index=True)
dates = np.array(dates)[ii]
return revs, dates
def get_rev_fiberassign_changes(svndir, rev, subdirs=None):
"""
Returns a Table() listing the changed fiberassign files for a given svn revision.
Args:
svndir: full path to a svn checkout (e.g. "$DESI_TARGET/fiberassign/tiles/trunk") (str)
rev: the revision (int)
subdirs (optional, defaults to None):
comma-separated list of subfolders in svndir (e.g. "000,001,080") (str)
if provided, cuts the list of fiberassign files on subdirs
Returns:
d: a Table() with: TILEID, FILE, DATE, REVISION
Notes:
This function queries the server ("svn log"), so do not run in parallel!
"""
# AR query svn log
start = time()
output = subprocess.check_output(
["svn", "log", "-r", str(rev), "-v", svndir], stderr=subprocess.DEVNULL
)
output = output.strip().decode()
dt = time() - start
# AR first get date
line_date = [
line
for line in output.split("\n")
if line[:2] in ["r{}".format(i) for i in range(1, 10)]
][0]
date = line_date.split("|")[2].split()[0]
# AR now cut on lines related to fiberassign changes
lines = [
line
for line in output.split("\n")
if "/tiles/trunk" in line and "/fiberassign-" in line
]
changes = np.array([line.split()[0] for line in lines])
fns = np.array([line.split()[1] for line in lines])
#
myd = {key: [] for key in ["TILEID", "FILE", "DATE", "REVISION", "CHANGE"]}
dt = time() - start
# AR cut on subdirs
tmp_tileids = np.array([int(os.path.basename(fn)[12:18]) for fn in fns])
tmp_subdirs = np.array(["{:06d}".format(tileid)[:3] for tileid in tmp_tileids])
sel = np.isin(tmp_subdirs, subdirs.split(","))
fns, changes = fns[sel], changes[sel]
log.info("rev={}\tdate={}\tnfile={}\t(dt={:.1f}s)".format(rev, date, len(fns), dt))
# AR store
# AR (forcing types to protect cases where len(fns)=0)
d = Table()
d.meta["SVNDIR"] = svndir
d.meta["SUBDIRS"] = subdirs
d["TILEID"] = np.array([int(os.path.basename(fn)[12:18]) for fn in fns], dtype=int)
d["FILE"] = np.array(fns, dtype=str)
d["DATE"] = np.array([date for fn in fns], dtype="<U10")
d["REVISION"] = np.array([rev for fn in fns], dtype=int)
d["CHANGE"] = np.array(changes, dtype=str)
return d
def get_obstheta_corr(decs, has, clip_arcsec=600.):
"""
Returns the computed correction to be applied to the field rotation
computed in fiberassign.tiles.load_tiles().
The correction should be applied as: obsthetas[deg] -= obstheta_corrs[deg].
Args:
decs: tile declinations (float or np array of floats)
has: hour angles (float or np array of floats)
clip_arcsec (optional, defaults to 600): abs(obstheta_corrs) is
forced to be <obstheta_corrs (float)
Returns:
obstheta_corrs: correction (in deg) to be applied (float or np array of floats)
Notes:
See DocDB-8931 for details.
During observations, PlateMaker computes the required field rotation,
then asks the hexapod to be rotated by
ROTOFFST = FIELDROT - PM_REQ_FIELDROT,
where FIELDROT is the field rotation coming from fiberassign.tiles.load_tiles().
When abs(ROTOFFST)>600 arcsec, the move is denied, and the exposure aborted.
Correction computed here is a fit to ROTOFFST=f(DEC), plus a fit on the residuals.
"""
assert clip_arcsec >= 0
isoneval = isinstance(decs, float)
if isoneval:
decs, has = np.atleast_1d(decs), np.atleast_1d(has)
# AR fitted function, ROTOFFST[arcsec] = f(DEC)
# AR rescale decs into [0, 1]
xs = (90. - decs) / 180.
rotoffsts = 937.60578 - 697.06513 * xs ** -0.18835
# AR fitted function to the residuals, residuals[arcsec] = f(HA)
xs = (90. + has) / 180.
residuals = -113.90162 + 222.18009 * xs
sel = xs > 0.5
residuals[sel] = -245.49007 + 485.35700 * xs[sel]
# AR total correction
obstheta_corrs = rotoffsts + residuals
# AR clip
obstheta_corrs = np.clip(obstheta_corrs, -clip_arcsec, clip_arcsec)
# AR switch to degrees
obstheta_corrs /= 3600
if isoneval:
decs, has = decs[0], has[0]
return obstheta_corrs[0]
else:
return obstheta_corrs
def get_main_dtver(tileid, svndir=None):
"""
Retrieve the desitarget catalog version for a main tile.
Args:
tileid: tileid (int)
svndir (optional, defaults to $DESI_TARGET/fiberassign/tiles/trunk): svn folder (str)
Returns:
dtver: the desitarget catalog version (str)
Notes:
The information is obtained from the FAARGS keyword in the
zero-th extension of the fiberassign-TILEID.fits.gz file
in the svn folder.
As of Aug. 2025, the possible outputs are:
'-', '1.0.0', '1.1.1', '2.2.0', '3.0.0', '3.2.0'
The '-' output is returned if no such file exists.
"""
# AR default svn folder
if svndir is None:
svndir = os.path.join(os.getenv("DESI_TARGET"), "fiberassign", "tiles", "trunk")
tileidpad = "{:06d}".format(tileid)
fn = os.path.join(svndir, tileidpad[:3], "fiberassign-{}.fits.gz".format(tileidpad))
if os.path.isfile(fn):
faargs = fitsio.read_header(fn, 0)["FAARGS"].split()
return [faargs[i+1] for i in range(len(faargs)-1) if faargs[i] == "--dtver"][0]
else:
return "-"
def get_main_early_nonrepro_tiles():
"""
Table with the 157 early main tiles, designed with DTVER=1.0.0, and non-reproducible.
Args:
None
Returns:
d: Table array with ['TILEID', 'PASS', 'RA', 'DEC', 'PROGRAM', 'IN_DESI', 'EBV_MED']
Notes:
The table is stored in data/tiles-main-dtver-1.0.0.ecsv.
It has been generated with:
d = Table.read("/global/cfs/cdirs/desi/survey/ops/surveyops/trunk/ops/tiles-main.ecsv")
pool = multiprocessing.Pool(processes=256)
with pool:
dtvers = np.array(pool.map(get_main_dtver, d["TILEID"]))
sel = dtvers == "1.0.0"
d = d[sel]["TILEID", "PASS", "RA", "DEC", "PROGRAM", "IN_DESI", "EBV_MED"]
d.write("tiles-main-dtver-1.0.0.ecsv")
"""
fn = resources.files("fiberassign") / "data" / "tiles-main-dtver-1.0.0.ecsv"
d = Table.read(fn)
return d