changeable cache dir

thomaswmorris · Jan 25, 2025 · 16f6822 · 16f6822
1 parent 4066561
commit 16f6822
Show file tree

Hide file tree

Showing 9 changed files with 107 additions and 48 deletions.
diff --git a/maria/__init__.py b/maria/__init__.py
@@ -12,7 +12,7 @@
 from .plan import Plan, all_plans, get_plan  # noqa
 from .sim import Simulation  # noqa
 from .site import Site, all_regions, all_sites, get_site  # noqa
-from .io import fetch  # noqa
+from .io import fetch, set_cache_dir  # noqa
 
 logging.basicConfig(
     level=logging.INFO,

diff --git a/maria/_version.py b/maria/_version.py
@@ -1,12 +1,10 @@
 # file generated by setuptools_scm
 # don't change, don't track in version control
-from __future__ import annotations
-
 TYPE_CHECKING = False
 if TYPE_CHECKING:
     from typing import Tuple, Union
 
-    VERSION_TUPLE = tuple[Union[int, str], ...]
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
 
@@ -15,5 +13,5 @@
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = "1.0.1.dev50"
-__version_tuple__ = version_tuple = (1, 0, 1, "dev50")
+__version__ = version = "1.0.1.dev65"
+__version_tuple__ = version_tuple = (1, 0, 1, "dev65")
diff --git a/maria/coords/coordinates.py b/maria/coords/coordinates.py
@@ -335,9 +335,11 @@ def compute_points(self):
 
     def center(self, frame=None):
         frame = frame or self.frame
-        return get_center_phi_theta(
-            getattr(self, frames[frame]["phi"]),
-            getattr(self, frames[frame]["theta"]),
+        return np.array(
+            get_center_phi_theta(
+                getattr(self, frames[frame]["phi"]),
+                getattr(self, frames[frame]["theta"]),
+            )
         )
 
     def broadcast(self, offsets, frame, axis=0):

diff --git a/maria/io/caching.py b/maria/io/caching.py
@@ -9,6 +9,13 @@
 from ..utils.io import test_file
 
 logger = logging.getLogger("maria")
+here, this_filename = os.path.split(__file__)
+
+os.environ["MARIA_CACHE_DIR"] = "/tmp/maria-data"
+
+
+def set_cache_dir(directory):
+    os.environ["MARIA_CACHE_DIR"] = directory
 
 
 def cache_status(path: str, max_age: float = 30 * 86400, refresh: bool = False):
@@ -94,7 +101,7 @@ def fetch(
     """
     Fetch a file from the repo.
     """
-    cache_path = cache_path or f"/tmp/maria-data/{source_path}"
+    cache_path = cache_path or f"{os.environ['MARIA_CACHE_DIR']}/data/{source_path}"
     url = f"{url_base}/{source_path}"
 
     status = cache_status(cache_path, max_age=max_age, refresh=refresh)

diff --git a/maria/map/projected.py b/maria/map/projected.py
@@ -433,6 +433,8 @@ def plot(
                 if i_nu == 0:
                     ax.set_ylabel(rf"$\Delta\,\theta_y$ [{y.units_short}.]")
 
+                ax.set_aspect("equal")
+
         dummy_map = mpl.cm.ScalarMappable(
             mpl.colors.Normalize(vmin=vmin, vmax=vmax),
             cmap=cmap,

diff --git a/maria/sim/base.py b/maria/sim/base.py
@@ -3,6 +3,7 @@
 import arrow
 import logging
 import os
+import gc
 
 import numpy as np
 import time as ttime
@@ -189,6 +190,8 @@ def run(self):
             metadata=metadata,
         )
 
+        gc.collect()
+
         return tod
 
     def plot_counts(self, x_bins=100, y_bins=100):

diff --git a/maria/site/regions.csv b/maria/site/regions.csv
@@ -1,26 +1,26 @@
 ,location,country,notes,latitude,longitude,altitude,timezone,training_start,training_end,training_years
-boolardy,"Boolardy, Western Australia",Australia,"MRO, SKA",-26.697,116.631,395,Australia/Perth,2010-01-01,2024-10-01,14.7
-boston,"Cambridge, Massachusetts",USA,"Harvard University, MIT",42.374,-71.111,8,America/New_York,2015-01-01,2024-10-01,9.7
-chajnantor,"Llano de Chajnantor, Antofagasta",Chile,"ACT, ALMA, APEX, SO, TAO",-23.006,-67.759,5040,America/Santiago,1959-08-17,2024-10-01,64.6
-chiang_mai,"Chiang Mai, Thailand",Thailand,TNRT,18.864,99.217,395,Asia/Bangkok,2022-01-01,2024-10-01,2.0
-effelsberg,"Effelsberg, North Rhine-Westphalia",Germany,ERT,50.524,6.883,319,Europe/Berlin,2015-01-01,2024-10-01,9.7
-green_bank,"Green Bank, West Virginia",USA,Green Bank Observatory,38.433,-79.84,807,America/New_York,2000-01-01,2024-10-01,24.7
-mauna_kea,"Mauna Kea, Hawaii",USA,Mauna Kea Observatory,19.823,-155.475,4205,Pacific/Honolulu,2000-01-01,2024-10-01,24.7
-meerkat,"Meerkat National Park, Northern Cape",South Africa,"HERA, MeerKAT, SKA",-30.713,21.443,1075,Africa/Johannesburg,2010-01-01,2024-10-01,14.7
-metsahovi,"Metsähovi, Kirkkonummi",Finland,Metsähovi Radio Observatory,60.218,24.394,90,Europe/Helsinki,2018-01-01,2024-10-01,6.0
-minamimaki,"Minamimaki, Nagano",Japan,Nobeyama Observatory,35.942,138.476,1350,Asia/Tokyo,2015-01-01,2024-10-01,9.7
-mount_graham,"Mount Graham, Arizona",USA,"LBT, VATT",32.702,-109.89,3178,America/Phoenix,2015-01-08,2024-10-01,9.4
-narrabri,"Narrabri, New South Wales",Australia,ATCA,-30.314,149.562,237,Australia/Sydney,2015-01-01,2024-10-01,9.7
-ngari,"Ngari Prefecture, Tibet",China,AliCPT,32.326,80.026,5176,Asia/Shanghai,2010-01-01,2024-10-01,14.7
-owens_valley,"Owens Valley, California",USA,Owens Valley Radio Observatory,37.232,-118.295,1222,America/Los_Angeles,2010-01-01,2024-10-01,14.7
-pic_de_bure,"Plateau de Bure, Haute-Alpes",France,NOEMA,44.634,5.908,2552,Europe/Paris,2015-01-01,2024-10-01,9.7
-pico_veleta,"Pico Veleta, Granada",Spain,IRAM,37.066,-3.393,2850,Europe/Madrid,2015-01-01,2024-10-01,9.7
-princeton,"Princeton, New Jersey",USA,Princeton University,40.345,-74.651,58,America/New_York,2015-01-01,2024-10-01,9.7
-qitai,"Qitai, Xinjiang",China,QTT,43.601,89.699,1780,Asia/Urumqi,2015-01-01,2024-10-01,9.7
-san_agustin,"Plains of San Agustín, New Mexico",USA,VLA,34.079,-107.618,2120,America/Denver,2010-01-01,2024-10-01,14.7
-san_basilio,"San Basilio, Sardinia",Italy,"SRT, MISTRAL",39.493,9.245,600,Europe/Rome,2010-01-01,2024-10-01,14.7
-sierra_negra,"Sierra Negra, Puebla",Mexico,"HAWC, LMT",18.986,-97.314,4640,America/Mexico_City,2010-01-03,2024-10-01,14.7
-south_pole,South Pole,Antarctica,"BICEP2, DASI, GML, Keck, SPT",-90.0,0.0,2835,Antarctica/South_Pole,1996-06-22,2024-10-01,26.9
-summit_camp,"Summit Station, Greenland",Greenland,Greenland Telescope,72.583,-38.455,3202,America/Nuuk,2005-07-15,2024-10-01,19.2
-teide,"Mount Teide, Tenerife",Spain,Teide Observatory,28.301,-16.511,2390,Atlantic/Canary,2010-01-01,2024-10-01,14.7
-thule,Thule Air Base,Greenland,Thule Air Base,76.535,-68.686,77,America/Thule,2023-03-15,2024-10-01,1.5
+boolardy,"Boolardy, Western Australia",Australia,"MRO, SKA",-26.697,116.631,395,Australia/Perth,2010-01-01,2024-10-30,14.8
+boston,"Cambridge, Massachusetts",USA,"Harvard University, MIT",42.374,-71.111,8,America/New_York,2015-01-01,2024-10-30,9.8
+chajnantor,"Llano de Chajnantor, Antofagasta",Chile,"ACT, ALMA, APEX, SO, TAO",-23.006,-67.759,5040,America/Santiago,1959-08-17,2024-10-30,64.7
+chiang_mai,"Chiang Mai, Thailand",Thailand,TNRT,18.864,99.217,395,Asia/Bangkok,2022-01-01,2024-10-30,2.1
+effelsberg,"Effelsberg, North Rhine-Westphalia",Germany,ERT,50.524,6.883,319,Europe/Berlin,2015-01-01,2024-10-30,9.8
+green_bank,"Green Bank, West Virginia",USA,Green Bank Observatory,38.433,-79.84,807,America/New_York,2000-01-01,2024-10-30,24.8
+mauna_kea,"Mauna Kea, Hawaii",USA,Mauna Kea Observatory,19.823,-155.475,4205,Pacific/Honolulu,2000-01-01,2024-10-30,24.8
+meerkat,"Meerkat National Park, Northern Cape",South Africa,"HERA, MeerKAT, SKA",-30.713,21.443,1075,Africa/Johannesburg,2010-01-01,2024-10-30,14.8
+metsahovi,"Metsähovi, Kirkkonummi",Finland,Metsähovi Radio Observatory,60.218,24.394,90,Europe/Helsinki,2018-01-01,2024-10-30,6.1
+minamimaki,"Minamimaki, Nagano",Japan,Nobeyama Observatory,35.942,138.476,1350,Asia/Tokyo,2015-01-01,2024-10-30,9.8
+mount_graham,"Mount Graham, Arizona",USA,"LBT, VATT",32.702,-109.89,3178,America/Phoenix,2015-01-08,2024-10-30,9.5
+narrabri,"Narrabri, New South Wales",Australia,ATCA,-30.314,149.562,237,Australia/Sydney,2015-01-01,2024-10-30,9.8
+ngari,"Ngari Prefecture, Tibet",China,AliCPT,32.326,80.026,5176,Asia/Shanghai,2010-01-01,2024-10-30,14.8
+owens_valley,"Owens Valley, California",USA,Owens Valley Radio Observatory,37.232,-118.295,1222,America/Los_Angeles,2010-01-01,2024-10-30,14.8
+pic_de_bure,"Plateau de Bure, Haute-Alpes",France,NOEMA,44.634,5.908,2552,Europe/Paris,2015-01-01,2024-10-30,9.8
+pico_veleta,"Pico Veleta, Granada",Spain,IRAM,37.066,-3.393,2850,Europe/Madrid,2015-01-01,2024-10-30,9.8
+princeton,"Princeton, New Jersey",USA,Princeton University,40.345,-74.651,58,America/New_York,2015-01-01,2024-10-30,9.8
+qitai,"Qitai, Xinjiang",China,QTT,43.601,89.699,1780,Asia/Urumqi,2015-01-01,2024-10-14,9.8
+san_agustin,"Plains of San Agustín, New Mexico",USA,VLA,34.079,-107.618,2120,America/Denver,2010-01-01,2024-10-30,14.8
+san_basilio,"San Basilio, Sardinia",Italy,"SRT, MISTRAL",39.493,9.245,600,Europe/Rome,2010-01-01,2024-10-30,14.8
+sierra_negra,"Sierra Negra, Puebla",Mexico,"HAWC, LMT",18.986,-97.314,4640,America/Mexico_City,2010-01-03,2024-10-26,14.8
+south_pole,South Pole,Antarctica,"BICEP2, DASI, GML, Keck, SPT",-90.0,0.0,2835,Antarctica/South_Pole,1996-06-22,2024-10-30,27.0
+summit_camp,"Summit Station, Greenland",Greenland,Greenland Telescope,72.583,-38.455,3202,America/Nuuk,2005-07-13,2024-10-30,19.2
+teide,"Mount Teide, Tenerife",Spain,Teide Observatory,28.301,-16.511,2390,Atlantic/Canary,2010-01-01,2024-10-30,14.8
+thule,Thule Air Base,Greenland,Thule Air Base,76.535,-68.686,77,America/Thule,2023-03-15,2024-10-30,1.6
diff --git a/maria/tod/tod.py b/maria/tod/tod.py
@@ -142,7 +142,7 @@ def fields(self) -> list:
 
     @property
     def signal(self) -> da.Array:
-        return sum([getattr(self, field) for field in self.fields])
+        return sum([self.data[field] for field in self.fields])
 
     @property
     def duration(self) -> float:

diff --git a/maria/weather/__init__.py b/maria/weather/__init__.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 import scipy as sp
+import pandas as pd
 
 from ..io import fetch, DEFAULT_TIME_FORMAT
 from ..site import InvalidRegionError, all_regions, supported_regions_table
@@ -71,7 +72,7 @@ def __init__(
         self.override = override
         self.source = source
 
-        time = time or arrow.now().to("utc")
+        time = time if time is not None else arrow.now().to("utc")
         self.time = arrow.get(time)
 
         self.cache_path = fetch(
@@ -90,6 +91,8 @@ def __init__(
         self.utc_day_hour = get_utc_day_hour(self.time.timestamp())
         self.utc_year_day = get_utc_year_day(self.time.timestamp())
 
+        self.data = {}
+
         self.region_quantiles = {}
 
         with h5py.File(self.cache_path, "r") as f:
@@ -124,7 +127,7 @@ def __init__(
                     (year_day_side, day_hour_side, self.quantile_levels),
                     data[YEAR_DAY_EDGE_INDEX, DAY_HOUR_EDGE_INDEX],
                 )((self.utc_year_day, self.utc_day_hour, self.quantiles.get(attr, 0.5)))
-                setattr(self, attr, y)
+                self.data[attr] = y
 
         wind_speed_correction_factor = self.wind_speed / np.sqrt(
             self.wind_east**2 + self.wind_north**2,
@@ -145,6 +148,13 @@ def __init__(
                 self.temperature, force_absolute_humidity
             )
 
+    def __getattr__(self, attr):
+
+        if attr in self.data:
+            return self.data[attr]
+
+        raise AttributeError()
+
     @property
     def absolute_humidity(self):
         return relative_to_absolute_humidity(self.temperature, self.humidity)
@@ -161,17 +171,53 @@ def wind_bearing(self):
     # def wind_speed(self):
     #     return np.sqrt(self.wind_east**2 + self.wind_north**2 + self.wind_vertical**2)
 
+    def layers(self):
+
+        df = pd.DataFrame(self.data)
+        df.insert(0, "altitude", df.geopotential.values / g)
+        df.insert(1, "pressure", 1e2 * self.pressure_levels)
+
+        df = df.loc[df.altitude > self.base_altitude]
+        df.index = np.arange(len(df))
+
+        absolute_humidity = relative_to_absolute_humidity(df.temperature, df.humidity)
+
+        df.insert(2, "h_bottom", 0.0)
+        df.insert(3, "h_top", 0.0)
+        df.insert(4, "h_thickness", 0.0)
+        df.insert(5, "total_water", 0.0)
+        df.insert(6, "absolute_humidity", absolute_humidity)
+
+        h_values = df.altitude.values
+        h_bins = np.array(
+            [
+                self.base_altitude,
+                *(h_values[:-1] + h_values[1:]) / 2,
+                h_values[-1] + 100,
+            ]
+        )
+
+        for layer_index, (h1, h2) in enumerate(zip(h_bins[:-1], h_bins[1:])):
+
+            h = np.linspace(h1, h2, 256)
+
+            absolute_humidity = sp.interpolate.interp1d(
+                self.altitude,
+                self.absolute_humidity,
+                bounds_error=False,
+                fill_value="extrapolate",
+            )(h)
+
+            df.loc[layer_index, ["h_bottom", "h_top", "h_thickness"]] = h1, h2, h2 - h1
+            df.loc[layer_index, "total_water"] = np.trapezoid(absolute_humidity, x=h)
+
+        df.loc[:, "altitude"] = (h_bins[:-1] + h_bins[1:]) / 2
+
+        return df
+
     @property
     def pwv(self):
-        altitude_samples = np.linspace(
-            self.base_altitude,
-            self.geopotential.max() / g,
-            1024,
-        )
-        return np.trapezoid(
-            np.interp(altitude_samples, self.geopotential / g, self.absolute_humidity),
-            x=altitude_samples,
-        )
+        return self.layers().total_water.sum()
 
     def __call__(self, altitude):
         res = {}
@@ -185,5 +231,6 @@ def __repr__(self):
         parts = []
         parts.append(f"region={repr(self.region)}")
         parts.append(f"time={self.local_time.format(DEFAULT_TIME_FORMAT)}")
-        parts.append(f"pwv={self.pwv:.03f}")
+        parts.append(f"elevation={int(self.base_altitude)}m")
+        parts.append(f"pwv={self.pwv:.03f}mm")
         return f"Weather({', '.join(parts)})"