Merge pull request #549 from jlmaurer/fix_imports_for_nisar

Move imports around to make raider-base cleaner

CHANGELOG.md

@@ -24,6 +24,7 @@ and uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 + Update the integration height for raytracing from 50 km to 80 km
 + Reinstate test 3 (slant proj and ray trace), remove unused calls with ZREF
 + Add buffer to W/E for ERA5
++ refactor imports to allow for a cleaner raider-base
 + Add buffer to HRES when downloading as with the other models
 + Refactor to pass a weather file directly to fetch
 + Update staged weather models to reflect update to aligned grid

test/test_HRRR_ztd.py

@@ -1,5 +1,4 @@
 import os
-import pytest
 import subprocess
 import shutil
 import glob

test/test_delayFcns.py

@@ -4,14 +4,29 @@
 import numpy as np
 import xarray as xr
 
-
+from pyproj import CRS, Transformer
 from test import TEST_DIR
 
 from RAiDER.delayFcns import getInterpolators
+from RAiDER.delay import transformPoints
 
 SCENARIO1_DIR = os.path.join(TEST_DIR, "scenario_1", "golden_data")
 
 
+@pytest.fixture
+def hrrr_proj():
+    lon0 = 262.5
+    lat0 = 38.5
+    lat1 = 38.5
+    lat2 = 38.5
+    x0 = 0
+    y0 = 0
+    earth_radius = 6371229
+    proj = CRS(f'+proj=lcc +lat_1={lat1} +lat_2={lat2} +lat_0={lat0} +lon_0={lon0} +x_0={x0} +y_0={y0} +a={earth_radius} +b={earth_radius} +units=m +no_defs')
+
+    return proj
+
+
 @pytest.fixture
 def wmdata():
     return xr.load_dataset(os.path.join(SCENARIO1_DIR, 'HRRR_tropo_20200101T120000_ztd.nc'))
@@ -29,3 +44,60 @@ def test_getInterpolators_2(wmdata, caplog):
     getInterpolators(ds, kind='pointwise')
     assert 'Weather model contains NaNs!' in caplog.text, 'No warning was raised!'
 
+
+def test_transformPoints():
+    lats = np.array([10, 20, 30, 45, 75, 80, 90])
+    lons = np.array([0,-180,180, 90,-90,-20, 10])
+    hts  = np.array([0,   0,  0,  0,  0,  0,  0])
+
+    epsg4326 = CRS.from_epsg(4326)
+    ecef = CRS.from_epsg(4978)
+
+    out = transformPoints(lats, lons, hts, epsg4326, ecef)
+    y,x,z=out[:,0], out[:,1], out[:,2]
+
+    T = Transformer.from_crs(4978, 4326)
+
+    test = T.transform(x,y,z)
+    assert np.allclose(test[0], lats)
+    assert np.allclose(test[1], lons)
+    assert np.allclose(test[2], hts)
+
+
+def test_transformPoints_2(hrrr_proj):
+    hrrr_proj = hrrr_proj
+    lats = np.array([40, 45, 55])
+    lons = np.array([-90,-90, -90])
+    hts  = np.array([0,   0,  0])
+
+    epsg4326 = CRS.from_epsg(4326)
+
+    out = transformPoints(lats, lons, hts, epsg4326, hrrr_proj)
+    y,x,z=out[:,0], out[:,1], out[:,2]
+
+    T = Transformer.from_crs(hrrr_proj, 4326)
+
+    test = T.transform(x,y,z)
+    assert np.allclose(test[0], lats)
+    assert np.allclose(test[1], lons)
+    assert np.allclose(test[2], hts)
+
+
+def test_transformPoints_3():
+    lats = np.array([0, 0, 0])
+    lons = np.array([0,-90, 180])
+    hts  = np.array([0,   0,  0])
+
+    epsg4326 = CRS.from_epsg(4326)
+    ecef = CRS.from_epsg(4978)
+
+    out = transformPoints(lats, lons, hts, epsg4326, ecef)
+    y,x,z=out[:,0], out[:,1], out[:,2]
+
+    assert np.allclose(x, [6378137., 0, -6378137])
+    assert np.allclose(y, [0, -6378137., 0])
+    assert np.allclose(z, [0, 0, 0])
+
+
+
+

test/test_intersect.py

@@ -1,7 +1,8 @@
 import pandas as pd
-import rasterio
+# import rasterio
 
 from scipy.interpolate import griddata
+import rasterio
 
 from test import *
 
@@ -13,10 +14,10 @@ def test_cube_intersect(wm):
     os.makedirs(SCENARIO_DIR, exist_ok=True)
 
     ## make the lat lon grid
-    S, N, W, E = 33.5, 34, -118.0, -117.5
+    # S, N, W, E = 33.5, 34, -118.0, -117.5
     date       = 20200130
     time       ='13:52:45'
-    f_lat, f_lon = makeLatLonGrid([S, N, W, E], 'LA', SCENARIO_DIR, 0.25)
+    # f_lat, f_lon = makeLatLonGrid([S, N, W, E], 'LA', SCENARIO_DIR, 0.25)
 
     ## make the template file
     grp = {
@@ -54,11 +55,8 @@ def test_cube_intersect(wm):
     lons = rasterio.open(lonf).read(1)
     hyd = griddata(np.stack([lons.flatten(), lats.flatten()], axis=-1), hyd.flatten(), (-100.6, 16.15), method='nearest')
 
-    # test for equality with golden data
     np.testing.assert_almost_equal(hyd, gold[wm], decimal=4)
 
-    return
-
 
 
 @pytest.mark.parametrize('wm', 'ERA5'.split())

test/test_synthetic.py

@@ -5,10 +5,9 @@
 from RAiDER.llreader import BoundingBox
 from RAiDER.models.weatherModel import make_weather_model_filename
 from RAiDER.losreader import Raytracing, build_ray
-from RAiDER.utilFcns import lla2ecef, ecef2lla
+from RAiDER.utilFcns import lla2ecef
 from RAiDER.cli.validators import modelName2Module
 
-from RAiDER.constants import _ZREF
 from test import *
 
 
@@ -61,7 +60,6 @@ def update_model(wm_file:str, wm_eq_type:str, wm_dir:str='weather_files_synth'):
 
     ds.close()
     del ds
-    print ('Wrote synthetic weather model file to:', dst)
     return dst
 
 

test/test_weather_model.py

@@ -15,7 +15,6 @@
     make_raw_weather_data_filename,
     make_weather_model_filename,
 )
-from RAiDER.llreader import BoundingBox
 from RAiDER.models.erai import ERAI
 from RAiDER.models.era5 import ERA5
 from RAiDER.models.era5t import ERA5T

tools/RAiDER/delay.py

@@ -12,9 +12,11 @@
 models are accessed as NETCDF files and should have "wet" "hydro"
 "wet_total" and "hydro_total" fields specified.
 """
+import os
 import pyproj
 import xarray
 
+from datetime import datetime
 from pyproj import CRS, Transformer
 from typing import List, Union
 
@@ -24,11 +26,6 @@
 from RAiDER.delayFcns import getInterpolators
 from RAiDER.logger import logger
 from RAiDER.losreader import build_ray
-from RAiDER.utilFcns import (
-    lla2ecef,  writeResultsToXarray,
-    rio_profile, transformPoints,
-)
-
 
 ###############################################################################
 def tropo_delay(
@@ -62,12 +59,12 @@ def tropo_delay(
     crs = CRS(out_proj)
 
     # Load CRS from weather model file
-    wm_proj = rio_profile(f"netcdf:{weather_model_file}:t")["crs"]
-    if wm_proj is None:
-       logger.warning("WARNING: I can't find a CRS in the weather model file, so I will assume you are using WGS84")
-       wm_proj = CRS.from_epsg(4326)
-    else:
-        wm_proj = CRS.from_wkt(wm_proj.to_wkt())
+    with xarray.load_dataset(weather_model_file) as ds:
+        try:
+            wm_proj = CRS.from_wkt(ds['proj'].attrs['crs_wkt'])
+        except KeyError:
+            logger.warning("WARNING: I can't find a CRS in the weather model file, so I will assume you are using WGS84")
+            wm_proj = CRS.from_epsg(4326)
 
     # get heights
     with xarray.load_dataset(weather_model_file) as ds:
@@ -228,6 +225,7 @@ def _build_cube_ray(
     """
     # Get model heights in an array
     # Assumption: All interpolators here are on the same grid
+    T = Transformer.from_crs(4326, 4978, always_xy=True)
     model_zs = interpolators[0].grid[2]
 
     # Create a regular 2D grid
@@ -242,10 +240,8 @@ def _build_cube_ray(
 
     # Various transformers needed here
     epsg4326 = CRS.from_epsg(4326)
-    cube_to_llh = Transformer.from_crs(pts_crs, epsg4326,
-                                       always_xy=True)
-    ecef_to_model = Transformer.from_crs(CRS.from_epsg(4978), model_crs,
-                                         always_xy=True)
+    cube_to_llh = Transformer.from_crs(pts_crs, epsg4326, always_xy=True)
+    ecef_to_model = Transformer.from_crs(CRS.from_epsg(4978), model_crs, always_xy=True)
 
     # Loop over heights of output cube and compute delays
     for hh, ht in enumerate(zpts):
@@ -259,7 +255,7 @@ def _build_cube_ray(
         else:
             llh = [xx, yy, np.full(yy.shape, ht)]
 
-        xyz = np.stack(lla2ecef(llh[1], llh[0], np.full(yy.shape, ht)), axis=-1)
+        xyz = np.stack(T.transform(llh[0], llh[1], llh[2]), axis=-1)
 
         # Step 2 - get LOS vectors for targets
         LOS = los.getLookVectors(ht, llh, xyz, yy)
@@ -324,3 +320,108 @@ def _build_cube_ray(
 
     if output_created_here:
         return outputArrs
+
+
+def writeResultsToXarray(dt, xpts, ypts, zpts, crs, wetDelay, hydroDelay, weather_model_file, out_type):
+    '''
+    write a 1-D array to a NETCDF5 file
+    '''
+       # Modify this as needed for NISAR / other projects
+    ds = xarray.Dataset(
+        data_vars=dict(
+            wet=(["z", "y", "x"],
+                 wetDelay,
+                 {"units" : "m",
+                  "description": f"wet {out_type} delay",
+                  # 'crs': crs.to_epsg(),
+                  "grid_mapping": "crs",
+
+                 }),
+            hydro=(["z", "y", "x"],
+                   hydroDelay,
+                   {"units": "m",
+                    # 'crs': crs.to_epsg(),
+                    "description": f"hydrostatic {out_type} delay",
+                    "grid_mapping": "crs",
+                   }),
+        ),
+        coords=dict(
+            x=(["x"], xpts),
+            y=(["y"], ypts),
+            z=(["z"], zpts),
+        ),
+        attrs=dict(
+            Conventions="CF-1.7",
+            title="RAiDER geo cube",
+            source=os.path.basename(weather_model_file),
+            history=str(datetime.utcnow()) + " RAiDER",
+            description=f"RAiDER geo cube - {out_type}",
+            reference_time=str(dt),
+        ),
+    )
+
+    # Write projection system mapping
+    ds["crs"] = int(-2147483647) # dummy placeholder
+    for k, v in crs.to_cf().items():
+        ds.crs.attrs[k] = v
+
+    # Write z-axis information
+    ds.z.attrs["axis"] = "Z"
+    ds.z.attrs["units"] = "m"
+    ds.z.attrs["description"] = "height above ellipsoid"
+
+    # If in degrees
+    if crs.axis_info[0].unit_name == "degree":
+        ds.y.attrs["units"] = "degrees_north"
+        ds.y.attrs["standard_name"] = "latitude"
+        ds.y.attrs["long_name"] = "latitude"
+
+        ds.x.attrs["units"] = "degrees_east"
+        ds.x.attrs["standard_name"] = "longitude"
+        ds.x.attrs["long_name"] = "longitude"
+
+    else:
+        ds.y.attrs["axis"] = "Y"
+        ds.y.attrs["standard_name"] = "projection_y_coordinate"
+        ds.y.attrs["long_name"] = "y-coordinate in projected coordinate system"
+        ds.y.attrs["units"] = "m"
+
+        ds.x.attrs["axis"] = "X"
+        ds.x.attrs["standard_name"] = "projection_x_coordinate"
+        ds.x.attrs["long_name"] = "x-coordinate in projected coordinate system"
+        ds.x.attrs["units"] = "m"
+
+    return ds
+
+
+
+def transformPoints(lats: np.ndarray, lons: np.ndarray, hgts: np.ndarray, old_proj: CRS, new_proj: CRS) -> np.ndarray:
+    '''
+    Transform lat/lon/hgt data to an array of points in a new
+    projection
+
+    Args:
+        lats: ndarray   - WGS-84 latitude (EPSG: 4326)
+        lons: ndarray   - ditto for longitude
+        hgts: ndarray   - Ellipsoidal height in meters
+        old_proj: CRS   - the original projection of the points
+        new_proj: CRS   - the new projection in which to return the points
+
+    Returns:
+        ndarray: the array of query points in the weather model coordinate system (YX)
+    '''
+    # Flags for flipping inputs or outputs
+    if not isinstance(new_proj, CRS):
+        new_proj = CRS.from_epsg(new_proj.lstrip('EPSG:'))
+    if not isinstance(old_proj, CRS):
+        old_proj = CRS.from_epsg(old_proj.lstrip('EPSG:'))
+
+    t = Transformer.from_crs(old_proj, new_proj, always_xy=True)
+
+    # in_flip = old_proj.axis_info[0].direction
+    # out_flip = new_proj.axis_info[0].direction
+
+    res  = t.transform(lons, lats, hgts)
+
+    # lat/lon/height
+    return  np.stack([res[1], res[0], res[2]], axis=-1)

tools/RAiDER/delayFcns.py

@@ -5,13 +5,17 @@
 # RESERVED. United States Government Sponsorship acknowledged.
 #
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-import multiprocessing as mp
+try:
+    import multiprocessing as mp
+except ImportError:
+    mp = None
+
 import xarray
 
 import numpy as np
+
 from scipy.interpolate import RegularGridInterpolator as Interpolator
 
-from RAiDER.utilFcns import transformPoints
 from RAiDER.logger import logger
 
 
@@ -50,7 +54,7 @@ def getInterpolators(wm_file, kind='pointwise', shared=False):
         wet   = make_shared_raw(wet)
         hydro = make_shared_raw(hydro)
 
-
+    
     ifWet = Interpolator((ys_wm, xs_wm, zs_wm), wet, fill_value=np.nan, bounds_error = False)
     ifHydro = Interpolator((ys_wm, xs_wm, zs_wm), hydro, fill_value=np.nan, bounds_error = False)
 
@@ -62,6 +66,9 @@ def make_shared_raw(inarr):
     Make numpy view array of mp.Array
     """
     # Create flat shared array
+    if mp is None:
+        raise ImportError('multiprocessing is not available')
+
     shared_arr = mp.RawArray('d', inarr.size)
     # Create a numpy view of it
     shared_arr_np = np.ndarray(inarr.shape, dtype=np.float64,