diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml new file mode 100644 index 0000000..6c1374e --- /dev/null +++ b/.github/workflows/python-package.yml @@ -0,0 +1,47 @@ +# This workflow will install Python dependencies, run tests and lint with a variety of Python versions +# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python + +name: Run pytests + +on: + push: + branches: [ "master" ] + pull_request: + branches: [ "master" ] + +jobs: + test: + + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install environment with micromamba + uses: mamba-org/setup-micromamba@v1 + with: + environment-file: tests/test-env.yml + create-args: python=${{ matrix.python-version }} + init-shell: bash + - name: Test with pytest + run: | + eval "$(micromamba shell hook --shell bash)" + micromamba activate pyxlma-tests + coverage run --source=pyxlma -m pytest --mpl --mpl-baseline-path=tests/truth/images/ --mpl-generate-summary=html,json --mpl-results-path=tests/mpl-results/ tests/ + coverage xml + - name: Upload coverage reports to Codecov + uses: codecov/codecov-action@v4 + env: + CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} + - name: Upload matplotlib test results + if: always() + uses: actions/upload-artifact@v4 + with: + name: matplotlib-results-${{ matrix.python-version }} + path: tests/mpl-results/ diff --git a/.gitignore b/.gitignore index ba097bc..39667ac 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,7 @@ coverage.xml *.cover .hypothesis/ .pytest_cache/ +tests/mpl-results/ # Translations *.mo diff --git a/README.md b/README.md index 1ec896c..e03ddf7 100644 --- a/README.md +++ b/README.md @@ -6,17 +6,52 @@ XLMA is a venerable IDL GUI program that diplays VHF Lightning Mapping Array dat Please use the issues tracker to discuss ideas and pull requests to contribute examples. # Installation -Clone this repostiory, switch to the development branch, and install. +Clone this repostiory install with pip. ``` git clone https://github.com/deeplycloudy/xlma-python.git cd xlma-python -git checkout cf-lma_format pip install -e . ``` Then, copy the `XLMA_plots.ipynb` notebook to wherever you'd like and start changing files, dates and times to show data from your case of interest. There also a notebook showing how to do flash sorting and save a new NetCDF file with those data. +# Dependencies +Required: + +- xarray (I/O requires the netcdf4 backend) +- pandas +- numpy + +Flash clustering: + +- scikit-learn +- scipy +- pyproj + +Plotting: + +- matplotlib +- cartopy +- metpy (optionally, for US county lines) + +GLM Plotting: +- glmtools (https://github.com/deeplycloudy/glmtools) + +Interactive: + +- jupyterlab (or, notebook) +- ipywidgets +- ipympl + +Building: + +- setuptools +- pytest-cov +- pytest-mpl +- lmatools (https://github.com/deeplycloudy/lmatools) +- ...and all of the above + # Technical architecture We envision a two-part package that keeps a clean separation between the core data model, analysis, and display. XLMA utilized a large, global `state` structure that stored all data, as well as the current data selection corresponding to the view in the GUI. Analysis then operated on whatever data was in the current selection. diff --git a/examples/LCFA_parallax_comparison.ipynb b/examples/LCFA_parallax_comparison.ipynb new file mode 100644 index 0000000..9d1cab3 --- /dev/null +++ b/examples/LCFA_parallax_comparison.ipynb @@ -0,0 +1,853 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from pyxlma.plot.xlma_base_plot import BlankPlot\n", + "from pyxlma.plot.xlma_plot_feature import plot_glm_events, plot_points, color_by_time\n", + "from glmtools.io.glm import GLMDataset, get_lutevents\n", + "from glmtools.io.lightning_ellipse import lightning_ellipse_rev, ltg_ellpse_rev\n", + "from glmtools.io.ccd import load_pixel_corner_lookup, quads_from_corner_lookup\n", + "from glmtools.io.imagery import get_goes_imager_proj\n", + "import datetime as dt\n", + "from cartopy import crs as ccrs\n", + "from pyxlma import coords\n", + "import numpy as np\n", + "import warnings\n", + "import xarray as xr\n", + "import pandas as pd\n", + "from matplotlib import pyplot as plt\n", + "warnings.filterwarnings(\"ignore\")" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
<xarray.Dataset> Size: 660kB\n", + "Dimensions: (number_of_events: 11094,\n", + " number_of_groups: 3821,\n", + " number_of_flashes: 212,\n", + " number_of_time_bounds: 2,\n", + " number_of_wavelength_bounds: 2,\n", + " number_of_field_of_view_bounds: 2)\n", + "Coordinates: (12/21)\n", + " event_id (number_of_events) uint32 44kB 16...\n", + " event_time_offset (number_of_events) datetime64[ns] 89kB ...\n", + " event_lat (number_of_events) float32 44kB -...\n", + " event_lon (number_of_events) float32 44kB -...\n", + " event_parent_group_id (number_of_events) uint32 44kB 66...\n", + " group_id (number_of_groups) uint32 15kB 66...\n", + " ... ...\n", + " product_time datetime64[ns] 8B 2023-12-24T00:5...\n", + " lightning_wavelength float32 4B ...\n", + " group_time_threshold float32 4B ...\n", + " flash_time_threshold float32 4B ...\n", + " lat_field_of_view float32 4B ...\n", + " lon_field_of_view float32 4B -75.0\n", + "Dimensions without coordinates: number_of_events, number_of_groups,\n", + " number_of_flashes, number_of_time_bounds,\n", + " number_of_wavelength_bounds,\n", + " number_of_field_of_view_bounds\n", + "Data variables: (12/37)\n", + " event_energy (number_of_events) float32 44kB 1...\n", + " group_frame_time_offset (number_of_groups) datetime64[ns] 31kB ...\n", + " group_area (number_of_groups) float32 15kB 2...\n", + " group_energy (number_of_groups) float32 15kB 1...\n", + " group_quality_flag (number_of_groups) float32 15kB ...\n", + " flash_frame_time_offset_of_first_event (number_of_flashes) datetime64[ns] 2kB ...\n", + " ... ...\n", + " event_x (number_of_events) float32 44kB 0...\n", + " event_y (number_of_events) float32 44kB -...\n", + " group_x (number_of_groups) float32 15kB 0...\n", + " group_y (number_of_groups) float32 15kB -...\n", + " flash_x (number_of_flashes) float32 848B ...\n", + " flash_y (number_of_flashes) float32 848B ...\n", + "Attributes: (12/29)\n", + " production_site: WCDAS\n", + " featureType: point\n", + " dataset_name: OR_GLM-L2-LCFA_G16_s20233580057000_e2023358005...\n", + " naming_authority: gov.nesdis.noaa\n", + " Conventions: CF-1.7\n", + " institution: DOC/NOAA/NESDIS> U.S. Department of Commerce, ...\n", + " ... ...\n", + " time_coverage_start: 2023-12-24T00:57:00.0Z\n", + " time_coverage_end: 2023-12-24T00:57:20.0Z\n", + " production_data_source: Realtime\n", + " production_environment: OE\n", + " LUT_Filenames: GLM_CALINR_AllFilters(FM1E_CDRL79RevK_DO_09_00...\n", + " id: a4140f6a-9068-4f60-a714-5cc0a5218f56
','active']) # Drop the station name column that has a redundant station letter code # as part of the name and join on station letter code. - station_combo = stations.set_index('ID').drop(columns=['Name']).join( - overview.set_index('ID')) + station_combo = stations.set_index(['ID','Name']).join( + overview.set_index(['ID','Name'])) self.stations = station_combo.reset_index(level=station_combo.index.names) def gen_sta_info(self): @@ -289,7 +431,7 @@ def gen_sta_info(self): of the station names. """ nstations = self.station_data_start-self.station_info_start-1 - with gzip.open(self.file) as f: + with open_gzip_or_dat(self.file) as f: for i in range(self.station_info_start+1): line = next(f) for line in range(nstations): @@ -307,7 +449,7 @@ def gen_sta_data(self): """ nstations = self.station_data_start-self.station_info_start-1 - with gzip.open(self.file) as f: + with open_gzip_or_dat(self.file) as f: for i in range(self.station_data_start+1): line = next(f) for line in range(nstations): @@ -332,8 +474,12 @@ def readfile(self): """ # Read in data. Catch case where there is no data. try: - lmad = pd.read_csv(self.file,compression='gzip',delim_whitespace=True, - header=None,skiprows=self.data_starts+1,error_bad_lines=False) + if self.file.endswith('.gz'): + comp = 'gzip' + else: + comp = None + lmad = pd.read_csv(self.file,compression=comp,sep='\\s+', + header=None,skiprows=self.data_starts+1,on_bad_lines='skip') lmad.columns = self.names except pd.errors.EmptyDataError: lmad = pd.DataFrame(columns=self.names) diff --git a/pyxlma/lmalib/lma_intercept_rhi.py b/pyxlma/lmalib/lma_intercept_rhi.py new file mode 100644 index 0000000..e7a810a --- /dev/null +++ b/pyxlma/lmalib/lma_intercept_rhi.py @@ -0,0 +1,308 @@ +from pyxlma.coords import RadarCoordinateSystem, TangentPlaneCartesianSystem, GeographicSystem +import numpy as np +import datetime as dt +import pandas as pd + + +def rcs_to_tps(radar_latitude, radar_longitude, radar_altitude, radar_azimuth): + """ + Find the unit vector coordinates (east, north, up) of the plane of a radar RHI scan. + + Creates a azimuth, elevation, range and tangent plane cartesian system at the radar's latitude and longitude, + and converts the RHI azimuth direction to the tangent plane coordinate system. + + Parameters + ---------- + radar_latitude : `float` + Latitude of the radar in degrees. + radar_longitude : `float` + Longitude of the radar in degrees. + radar_altitude : `float` + Altitude of the radar in meters. + radar_azimuth : `float` + Azimuth of the RHI scan in degrees. + + Returns + ---------- + X : `numpy.ndarray` + A 1x2 array representing the start and end points eastward component of the RHI scan. + Y : `numpy.ndarray` + A 1x2 array representing the start and end points northward component of the RHI scan. + Z : `numpy.ndarray` + A 1x2 array representing the start and end points upward component of the RHI scan. + """ + + # Coordinates Systems + rcs = RadarCoordinateSystem(radar_latitude, radar_longitude, radar_altitude) + tps = TangentPlaneCartesianSystem(radar_latitude, radar_longitude, radar_altitude) + + # - Elevations, azimuth, range + r = np.array([0, 1]) + + els = np.array([0]) + els = np.tensordot(els, np.ones_like(r), axes=0) + + azi = np.array([radar_azimuth]) + az = np.tensordot(azi, np.ones_like(r), axes=0) + + a, b, c = rcs.toECEF(r,az,els) + abc = np.vstack((a,b,c)) + # ECEF to TPS + n = tps.toLocal(abc) + X = n[0,:] + Y = n[1,:] + Z = n[2,:] + + X = np.reshape(X, (1, 2)) + Y = np.reshape(Y, (1, 2)) + Z = np.reshape(Z, (1, 2)) + + return X, Y, Z + + +def geo_to_tps(event_longitude, event_latitude, event_altitude, tps_latitude, tps_longitude, tps_altitude): + """ + Convert the latitude, longitude, and altitude of LMA VHF sources to x/y/z distances (in meters) from an arbitrary latitude, longitude, and altitude. + + Creates a tangent plane cartesian system at the latitude, longitude, and altitude provided, and converts the LMA VHF sources to the tangent plane coordinate system. + + Parameters + ---------- + event_longitude : `xarray.Dataset` + A pyxlma dataset containing latitude, longitude, and altitude of LMA VHF sources. + tps_latitude : `float` + Latitude of the tangent plane in degrees. + tps_longitude : `float` + Longitude of the tangent plane in degrees. + tps_altitude : `float` + Altitude of the tangent plane in meters. + + Returns + ---------- + Xlma : `numpy.ndarray` + A 1xN array representing the eastward distance (in meters) of the tangent plane center to the LMA VHF sources. + Ylma : `numpy.ndarray` + A 1xN array representing the northward distance (in meters) of the tangent plane center to the LMA VHF sources. + Zlma : `numpy.ndarray` + A 1xN array representing the upward distance (in meters) of the tangent plane center to the LMA VHF sources. + """ + # GeographicSystem GEO - Lat, lon, alt + geo = GeographicSystem() + # Tangent Plane Cartesian System TPS - + tps = TangentPlaneCartesianSystem(tps_latitude, tps_longitude, tps_altitude) + + # GEO to TPS + + d, e, h = geo.toECEF(event_longitude, event_latitude, event_altitude) + + + deh = np.vstack((d,e,h)) + m = tps.toLocal(deh) + + Xlma = m[0] + Ylma = m[1] + Zlma = m[2] + + return Xlma,Ylma,Zlma + + +def ortho_proj_lma(event_longitude, event_latitude, event_altitude, radar_latitude, radar_longitude, radar_altitude, radar_azimuth): + """ + Convert the latitude, longitude, and altitude of LMA VHF sources to distance along, distance from, and height above the ground of a radar RHI scan. + + Creates tangent plane at the radar's location and converts the LMA VHF sources to the tangent plane coordinate system, then rotates the tangent plane in the direction of the scan azimuth. + + + Parameters + ---------- + lma_file : `xarray.Dataset` + A pyxlma dataset containing latitude, longitude, and altitude of N number of LMA VHF sources. + radar_latitude : `float` + Latitude of the radar in degrees. + radar_longitude : `float` + Longitude of the radar in degrees. + radar_altitude : `float` + Altitude of the radar in meters. + radar_azimuth : `float` + Azimuth of the RHI scan in degrees. + + Returns + ---------- + lma_file_loc : `numpy.ndarray` + A Nx3 array representing the distance along, distance from, and height above the ground (in m) of the LMA VHF sources. + + """ + Xlma,Ylma,Zlma = geo_to_tps(event_longitude, event_latitude, event_altitude, radar_latitude, radar_longitude, radar_altitude) + X, Y, Z = rcs_to_tps(radar_latitude, radar_longitude, radar_altitude, radar_azimuth) + + lon_ini1 = X[0,0] + lat_ini1 = Y[0,0] + lon_fin1 = X[0,-1] + lat_fin1 = Y[0,-1] + + + dlon1 = lon_fin1 - lon_ini1 # dx + dlat1 = lat_fin1 - lat_ini1 # dy + ds1 = np.array((dlon1,dlat1)) + norm_ds1 = np.linalg.norm(ds1) + cross_ds1 = np.tensordot(ds1, ds1, (0,0)) + + # LMA + lma_file_n = np.column_stack((Xlma, Ylma)) + + lma_file_loc_par = np.zeros(shape=(len(lma_file_n), 2)) + lma_file_loc_perp = np.zeros(shape=(len(lma_file_n), 2)) + lma_file_loc = np.zeros(shape=(len(lma_file_n), 3)) + + # + # ################################## + # + # (Xlma[i],Ylma[i]).ds1 .ds1 + # ---------------------- + # ds1 . ds1 + # + # ################################## + # + lma_file_loc_tensor_x = np.tensordot(ds1,lma_file_n,(0,1)) + lma_file_loc_par = np.tensordot((lma_file_loc_tensor_x / cross_ds1 ),ds1,0) + + # + # ####################################################################### + # + # (Xlma[i],Ylma[i]) _ (Xlma[i],Ylma[i]).ds1 .ds1 + # ---------------------- + # ds1 . ds1 + # + ########################################################################## + # + lma_file_loc_perp = lma_file_n - lma_file_loc_par + + # + lma_file_loc[:,0] = np.sqrt(lma_file_loc_par[:,0]**2 + lma_file_loc_par[:,1]**2) + if radar_azimuth <= 90 or radar_azimuth >= 270: + lma_file_loc[:,0][lma_file_loc_par[:,1] < 0] = -lma_file_loc[:,0][lma_file_loc_par[:,1] < 0] + elif radar_azimuth >= 180 and radar_azimuth < 270: + lma_file_loc[:,0][lma_file_loc_par[:,0] > 0] = -lma_file_loc[:,0][lma_file_loc_par[:,0] > 0] + else: + lma_file_loc[:,0][lma_file_loc_par[:,0] < 0] = -lma_file_loc[:,0][lma_file_loc_par[:,0] < 0] + lma_file_loc[:,1] = np.sqrt(lma_file_loc_perp[:,0]**2 + lma_file_loc_perp[:,1]**2) + lma_file_loc[:,2] = Zlma + + return lma_file_loc + + +def find_points_near_rhi(event_longitude, event_latitude, event_altitude, event_time, + radar_latitude, radar_longitude, radar_altitude, radar_azimuth, radar_scan_time, + distance_threshold=1000, time_threshold=30): + """ + Find the LMA VHF sources near a radar RHI scan. + + Creates tangent plane at the radar's location and converts the LMA VHF sources to the tangent plane coordinate system, then rotates the tangent plane in the direction of the scan azimuth. + Filters RHI scan points based on distance and time thresholds. + + + Parameters + ---------- + event_longitude : array-like + An array of the latitudes of events to be transformed. + event_latitude : array-like + An array of the latitudes of events to be transformed. + event_altitude : array-like + An array of the altitudes of events to be transformed. + event_time : array-like + An array of the times of events to be transformed. + radar_latitude : `float` + Latitude of the radar in degrees. + radar_longitude : `float` + Longitude of the radar in degrees. + radar_altitude : `float` + Altitude of the radar in meters. + radar_azimuth : `float` + Azimuth of the RHI scan in degrees. + radar_scan_time : `datetime.datetime` or `numpy.datetime64` or `pandas.Timestamp` + Time of the RHI scan. + distance_threshold : `float` + Maximum distance from the radar to the LMA VHF sources in meters. Default is 1000. + time_threshold : `float` + Number of seconds before and after the RHI scan time to include LMA VHF sources. Default is 30. (total length: 1 minute) + + + Returns + ---------- + lma_range : `numpy.ndarray` + A 1D array representing the distance along the tangent plane in the direction of the RHI scan. + lma_dist : `numpy.ndarray` + A 1D array representing the distance from the radar RHI scan plane to each filtered LMA point. + lma_alt : `numpy.ndarray` + A 1D array representing the height above the tangent plane centered at radar level of each filtered LMA point. + point_mask : `numpy.ndarray` + A 1D array of booleans representing the VHF points that were included in the return. + """ + + radar_azimuth = radar_azimuth % 360 + + radar_scan_time = np.array([radar_scan_time]).astype('datetime64[s]').astype(dt.datetime) + + projected_lma = ortho_proj_lma(event_longitude, event_latitude, event_altitude, + radar_latitude, radar_longitude, radar_altitude, radar_azimuth) + lma_range = projected_lma[:,0] + lma_dist = projected_lma[:,1] + lma_alt = projected_lma[:,2] + + if isinstance(event_time, pd.Series): + event_time = event_time.values + lma_times = event_time.astype('datetime64[s]').astype(dt.datetime) + point_mask = np.ones_like(lma_range, dtype=bool) + if distance_threshold is not None: + point_mask = np.logical_and(point_mask, lma_dist < distance_threshold) + if time_threshold is not None: + point_mask = np.logical_and(point_mask, + np.abs(lma_times - radar_scan_time) < dt.timedelta(seconds=time_threshold)) + lma_range = lma_range[point_mask] + lma_dist = lma_dist[point_mask] + lma_alt = lma_alt[point_mask] + return lma_range, lma_dist, lma_alt, point_mask + + +def find_lma_points_near_rhi(lma_file, radar_latitude, radar_longitude, radar_altitude, radar_azimuth, radar_scan_time, distance_threshold=1000, time_threshold=30): + """ + Find the LMA VHF sources near a radar RHI scan. + + Creates tangent plane at the radar's location and converts the LMA VHF sources to the tangent plane coordinate system, then rotates the tangent plane in the direction of the scan azimuth. + Filters RHI scan points based on distance and time thresholds. + + + Parameters + ---------- + lma_file : `xarray.Dataset` + A pyxlma dataset containing latitude, longitude, and altitude, and event_id of N number of LMA VHF sources. + radar_latitude : `float` + Latitude of the radar in degrees. + radar_longitude : `float` + Longitude of the radar in degrees. + radar_altitude : `float` + Altitude of the radar in meters. + radar_azimuth : `float` + Azimuth of the RHI scan in degrees. + radar_scan_time : `datetime.datetime` or `numpy.datetime64` or `pandas.Timestamp` + Time of the RHI scan. + distance_threshold : `float` + Maximum distance from the radar to the LMA VHF sources in meters. Default is 1000. + time_threshold : `float` + Number of seconds before and after the RHI scan time to include LMA VHF sources. Default is 30. (total length: 1 minute) + + + Returns + ---------- + lma_range : `numpy.ndarray` + A 1D array representing the distance along the tangent plane in the direction of the RHI scan. + lma_dist : `numpy.ndarray` + A 1D array representing the distance from the radar RHI scan plane to each filtered LMA point. + lma_alt : `numpy.ndarray` + A 1D array representing the height above the tangent plane centered at radar level of each filtered LMA point. + point_mask : `numpy.ndarray` + A 1D array of booleans representing the VHF points that were included in the return. + """ + return find_points_near_rhi(lma_file.event_longitude.data, lma_file.event_latitude.data, lma_file.event_altitude.data, + lma_file.event_time.data, radar_latitude, radar_longitude, radar_altitude, radar_azimuth, + radar_scan_time, distance_threshold, time_threshold) diff --git a/pyxlma/lmalib/traversal.py b/pyxlma/lmalib/traversal.py index 1f84ddf..54cdc01 100644 --- a/pyxlma/lmalib/traversal.py +++ b/pyxlma/lmalib/traversal.py @@ -77,7 +77,7 @@ def __init__(self, dataset, entity_id_vars, parent_id_vars): self.child_to_parent = collections.OrderedDict() self.parent_to_child = collections.OrderedDict() for (entity_var, parent_var) in self._descend(): - if dataset.dims[dataset[entity_var].dims[0]] == 0: + if dataset.sizes[dataset[entity_var].dims[0]] == 0: # No data, so groupby will fail in xarray > 0.13 entity_grouper = None else: @@ -86,7 +86,7 @@ def __init__(self, dataset, entity_id_vars, parent_id_vars): if parent_var is None: parent_grouper = None else: - if dataset.dims[dataset[parent_var].dims[0]] == 0: + if dataset.sizes[dataset[parent_var].dims[0]] == 0: # No data, so groupby will fail in xarray > 0.13 parent_grouper = None else: @@ -219,7 +219,7 @@ def reduce_to_entities(self, entity_id_var, entity_ids): for e_var, p_var in self._descend(): if prune == True: p_group = self.parent_groups[p_var].groups - e_iter = (p_group[eid] for eid in last_entity_ids + e_iter = (np.atleast_1d(p_group[eid]) for eid in last_entity_ids if eid in p_group) e_idx = list(itertools.chain.from_iterable(e_iter)) if len(e_idx) == 0: @@ -239,7 +239,7 @@ def reduce_to_entities(self, entity_id_var, entity_ids): # through. prune = True e_group = self.entity_groups[e_var].groups - e_iter = (e_group[eid] for eid in entity_ids + e_iter = (np.atleast_1d(e_group[eid]) for eid in entity_ids if eid in e_group) e_idx = list(itertools.chain.from_iterable(e_iter)) last_entity_ids = entity_ids # == dataset[e_var].data @@ -263,7 +263,7 @@ def reduce_to_entities(self, entity_id_var, entity_ids): for e_var, p_var in self._ascend(): if (prune == True): e_group = self.entity_groups[e_var].groups - e_iter = (e_group[eid] for eid in last_entity_ids + e_iter = (np.atleast_1d(e_group[eid]) for eid in last_entity_ids if eid in e_group) e_idx = list(itertools.chain.from_iterable(e_iter)) if len(e_idx) == 0: diff --git a/pyxlma/plot/interactive.py b/pyxlma/plot/interactive.py index 61ebfce..c8038d0 100644 --- a/pyxlma/plot/interactive.py +++ b/pyxlma/plot/interactive.py @@ -264,6 +264,14 @@ def make_plot(self): lon_set, lat_set, alt_set, time_set, selection = subset( lon_data, lat_data, alt_data, time_data, chi_data, station_data, xlim, ylim, zlim, tlim_sub, xchi, stationmin) + # Retain the current LMA data so that subclasses can access the current LMA + # data and compare to other data to be plotted. For instance, to calculate + # the offset between the time of the first LMA point and a ground strike point. + self.this_lma_lon = lon_set + self.this_lma_lat = lat_set + self.this_lma_alt = alt_set + self.this_lma_time = time_set + self.this_lma_sel = selection # if self.lma_plot is not None: # fig = self.lma_plot.fig diff --git a/pyxlma/plot/leader_speed.py b/pyxlma/plot/leader_speed.py new file mode 100644 index 0000000..8bc853d --- /dev/null +++ b/pyxlma/plot/leader_speed.py @@ -0,0 +1,91 @@ +import numpy as np +from numpy import radians, cos, sin, arcsin, sqrt +import matplotlib.pyplot as plt +import pandas as pd + +def haversine(lat0, lon0, lat, lon): + + R = 6378.137e3 # this is in meters. For Earth radius in kilometers use 6372.8 km + + dLat = radians(lat - lat0) + dLon = radians(lon - lon0) + lat1 = radians(lat0) + lat2 = radians(lat) + + a = sin(dLat/2)**2 + cos(lat1)*cos(lat2)*sin(dLon/2)**2 + c = 2*arcsin(sqrt(a)) + + return R * c + + +def get_time_distance(lat, lon, time, lat0, lon0, time0): + lat0=lat0*np.ones_like(lat) + lon0=lon0*np.ones_like(lon) + dt = time-time0 + + distance_from_origin = (haversine(lat0,lon0,lat,lon)) + if(type(dt) == pd.core.series.Series): + dt = dt.to_numpy() + + time_from_origin = ((dt).astype('timedelta64[ns]').astype(float)/1e9) + return distance_from_origin, time_from_origin + + +def time_distance_plot_interactive(interactive_lma, ax): + lat = interactive_lma.this_lma_lat + lon = interactive_lma.this_lma_lon + alt = interactive_lma.this_lma_alt + time = interactive_lma.this_lma_time + first = np.nanargmin(time) + + distance_from_origin, time_from_origin = get_time_distance(lat, lon, time, + lat[first], lon[first], time[first]) + + art_out = time_distance_plot(ax, time_from_origin, distance_from_origin) + + return art_out + + +def time_distance_plot(ax, time, distance, pad_sec = 10, **kwargs): + m_reference_lines = 10 + m = -m_reference_lines + dt = time[-1]-time[0] + x = np.linspace(0, dt + pad_sec, 100) + y = x * 2 * 10**4 + yy = x * 10**5 + yyy = x * 10**6 + + art_out = [] + + while (m < m_reference_lines): + art = ax.plot(x+(m/m_reference_lines), y, color = 'b') + art_out.extend(art) + art = ax.plot(x+(m/m_reference_lines), yy, color = 'r') + art_out.extend(art) + art = ax.plot(x+(m/m_reference_lines), yyy, color = 'g') + art_out.extend(art) + m+=1 + + art = ax.plot(x, y, color = 'b', label = 'Positive Leader') + art_out.extend(art) + art = ax.plot(x, yy, color = 'r', label = 'Negative Leader') + art_out.extend(art) + art = ax.plot(x, yyy, color = 'g', label = 'Dart Leader') + art_out.extend(art) + + ax.set_ylim(0, max(distance)+1000) + ax.set_xlim(-0.05, max(time)+0.1) + + sc = ax.scatter(time, distance, **kwargs) + art_out.append(sc) + #art = ax.legend(title = "Leader Type Key", loc = 2, fontsize=12, title_fontsize=16, framealpha=1) + #art_out.extend(art) + ax.set_title("Lightning Leader Speed", size=24) + ax.set_xlabel("Time from Origin (s)", size=12) + ax.set_ylabel("Distance from Origin (m)", size=12) + + return art_out + + + + \ No newline at end of file diff --git a/pyxlma/plot/xlma.py b/pyxlma/plot/xlma.py index 7947107..984ec27 100644 --- a/pyxlma/plot/xlma.py +++ b/pyxlma/plot/xlma.py @@ -2,6 +2,7 @@ import matplotlib.pyplot as plt import matplotlib.dates as md import datetime as dt +import pandas as pd from matplotlib.ticker import Formatter, FormatStrFormatter, MaxNLocator import cartopy @@ -145,7 +146,7 @@ def setup_figure(self, **kwargs): from pandas.plotting import register_matplotlib_converters register_matplotlib_converters() try: - self.datetime = self.data.event_time.to_dataframe().event_time + self.datetime = pd.Series(self.data.event_time.data.flatten()).dt.to_pydatetime() except: self.data_exists=False self.datetime = [self.stime,self.stime+dt.timedelta(minutes=10)] @@ -311,7 +312,7 @@ def histogram(self): plt.text(0.30, 0.80, '(c)', fontsize='x-large', weight='bold', horizontalalignment='center', verticalalignment='center', transform=self.ax_hist.transAxes) - + def plan_view(self): if self.data_exists == True: @@ -398,7 +399,7 @@ def inset_view(self, **kwargs): self.inset_size = kwargs['inset_size'] else: self.inset_size = 0.15 - self.inset = self.fig.add_axes([0.02, 0.01, 0.02+self.inset_size, + self.inset = self.fig.add_axes([0.02, 0.01, 0.02+self.inset_size, 0.01+self.inset_size],projection=ccrs.PlateCarree()) if self.data_exists==True: if self.readtype == 'lmatools': @@ -411,7 +412,7 @@ def inset_view(self, **kwargs): lon_data = self.data['event_longitude'][self.cond2].data lat_data = self.data['event_latitude'][self.cond2].data self.inset.hist2d(lon_data, lat_data, - bins=[int((self.xlim[1]+self.buffer*2 - self.xlim[0]) / xdiv), + bins=[int((self.xlim[1]+self.buffer*2 - self.xlim[0]) / xdiv), int((self.ylim[1]+self.buffer*2 - self.ylim[0]) / ydiv)], density=True, cmap=self.cmap, cmin=0.00001) @@ -421,7 +422,7 @@ def inset_view(self, **kwargs): self.inset.add_feature(COUNTIES, facecolor='none', edgecolor='gray') self.inset.add_feature(cfeature.BORDERS) self.inset.add_feature(cfeature.STATES.with_scale('10m')) - self.inset.set_extent([self.xlim[0]-self.buffer, self.xlim[1]+self.buffer, + self.inset.set_extent([self.xlim[0]-self.buffer, self.xlim[1]+self.buffer, self.ylim[0]-self.buffer, self.ylim[1]+self.buffer]) self.inset.plot([self.xlim[0],self.xlim[0],self.xlim[1],self.xlim[1],self.xlim[0]], [self.ylim[0],self.ylim[1],self.ylim[1],self.ylim[0],self.ylim[0]], diff --git a/pyxlma/plot/xlma_base_plot.py b/pyxlma/plot/xlma_base_plot.py index 8f3a465..c2f14b7 100644 --- a/pyxlma/plot/xlma_base_plot.py +++ b/pyxlma/plot/xlma_base_plot.py @@ -3,6 +3,7 @@ import matplotlib.dates as md import datetime as dt from matplotlib.ticker import Formatter, FormatStrFormatter, MaxNLocator +from matplotlib.dates import AutoDateLocator import cartopy import cartopy.crs as ccrs @@ -142,6 +143,13 @@ def plot(self, **kwargs): tfmt = '%H:%M:%S000' self.ax_th.set_xlim(self.tlim[0], self.tlim[1]) self.ax_th.xaxis.set_major_formatter(FractionalSecondFormatter(self.ax_th)) + # Importing pandas results in it overriding Matplotlib default AutoDateLocator, + # which in turn prevents more than one tick displaying for short time intervals. + # See this blog entry for details. + # https://notebook.community/azjps/matplotlib-tick-formatters/ notebooks/microsecond_precision At small time intervals, + # So, restore Matplotlib's AutoDateLocator, which as of version >3.3 correctly + # handles smaller time intervals. 3.2 did not! + self.ax_th.xaxis.set_major_locator(AutoDateLocator()) # Longitude-Altitue self.ax_lon.set_ylabel('Altitude (km MSL)') @@ -173,13 +181,12 @@ def plot(self, **kwargs): # Plan view if self.bkgmap==True: - self.ax_plan.add_feature(COUNTIES, facecolor='none', edgecolor='gray') + if COUNTIES != None: + self.ax_plan.add_feature(COUNTIES, facecolor='none', edgecolor='gray') self.ax_plan.add_feature(cfeature.BORDERS) self.ax_plan.add_feature(cfeature.STATES.with_scale('10m')) self.ax_plan.set_xlabel('Longitude (degrees)') self.ax_plan.set_ylabel('Latitude (degrees)') - self.ax_plan.set_xlim(self.xlim) - self.ax_plan.set_ylim(self.ylim) # lon_formatter = LongitudeFormatter(number_format='.2f', @@ -194,6 +201,9 @@ def plot(self, **kwargs): self.ax_plan.xaxis.set_major_formatter(self.majorFormatter) self.ax_plan.yaxis.set_major_formatter(self.majorFormatter) self.set_ax_plan_labels() + self.ax_plan.set_xlim(self.xlim) + self.ax_plan.set_ylim(self.ylim) + self.ax_plan.set_aspect('auto') def subplot_labels(plot): """ diff --git a/pyxlma/plot/xlma_plot_feature.py b/pyxlma/plot/xlma_plot_feature.py index e5dca43..1f01800 100644 --- a/pyxlma/plot/xlma_plot_feature.py +++ b/pyxlma/plot/xlma_plot_feature.py @@ -1,6 +1,9 @@ import matplotlib.pyplot as plt import numpy as np +import xarray as xr import matplotlib.dates as md +from matplotlib.patches import Polygon +from matplotlib.collections import PatchCollection def subset(lon_data, lat_data, alt_data, time_data, chi_data,station_data, @@ -25,19 +28,24 @@ def subset(lon_data, lat_data, alt_data, time_data, chi_data,station_data, return lon_data, lat_data, alt_data, time_data, selection -def color_by_time(time_array, tlim): +def color_by_time(time_array, tlim=None): """ - Generates colormap values for plotting VHF sources by time in a + Generates colormap values for plotting scatter points by time in a given time window Returns: min, max values, array by time """ - vmax = (tlim[1] - time_array.min()).total_seconds() - ldiff = time_array - time_array.min() - ldf = [] - for df in ldiff: - ldf.append(df.total_seconds()) - c = np.array(ldf) + if tlim is None: + tlim = np.array([np.atleast_1d(time_array.min())[0], + np.atleast_1d(time_array.max())[0]]) + nsToS = 1e9 + time_array = np.array(time_array).astype('datetime64[ns]').astype(float)/nsToS + tlim = np.atleast_1d(tlim).astype('datetime64[ns]').astype(float)/nsToS + + ldiff = time_array - tlim[0] + + vmax = tlim[1] - tlim[0] + c = ldiff.astype(float) vmin = 0 return vmin, vmax, c @@ -57,48 +65,249 @@ def setup_hist(lon_data, lat_data, alt_data, time_data, def plot_points(bk_plot, lon_data, lat_data, alt_data, time_data, - plot_cmap, plot_s, plot_vmin, plot_vmax, plot_c): + plot_cmap=None, plot_s=None, plot_vmin=None, plot_vmax=None, plot_c=None, edge_color='face', + edge_width=0, add_to_histogram=True, marker='o', **kwargs): """ Plot scatter points on an existing bk_plot object given x,y,z,t for each and defined plotting colormaps and ranges """ + + # before **kwargs was added to the function call, the following arguments + # were specified as keywords separately. This allows backwards compatibility: + if plot_cmap is None: + plot_cmap = kwargs.pop('cmap', plot_cmap) + if plot_s is None: + plot_s = kwargs.pop('s', plot_s) + if plot_vmin is None: + plot_vmin = kwargs.pop('vmin', plot_vmin) + if plot_vmax is None: + plot_vmax = kwargs.pop('vmax', plot_vmax) + if plot_c is None: + plot_c = kwargs.pop('c', plot_c) + if edge_color == 'face': + edge_color = kwargs.pop('edgecolors', edge_color) + if edge_width == 0: + edge_width = kwargs.pop('linewidths', edge_width) + art_plan = bk_plot.ax_plan.scatter(lon_data, lat_data, c=plot_c,vmin=plot_vmin, vmax=plot_vmax, cmap=plot_cmap, - s=plot_s,marker='o', edgecolors='none') + s=plot_s, marker=marker, linewidths=edge_width, edgecolors=edge_color, **kwargs) art_th = bk_plot.ax_th.scatter(time_data, alt_data, c=plot_c,vmin=plot_vmin, vmax=plot_vmax, cmap=plot_cmap, - s=plot_s,marker='o', edgecolors='none') + s=plot_s, marker=marker, linewidths=edge_width, edgecolors=edge_color, **kwargs) art_lon = bk_plot.ax_lon.scatter(lon_data, alt_data, c=plot_c,vmin=plot_vmin, vmax=plot_vmax, cmap=plot_cmap, - s=plot_s,marker='o', edgecolors='none') + s=plot_s, marker=marker, linewidths=edge_width, edgecolors=edge_color, **kwargs) art_lat = bk_plot.ax_lat.scatter(alt_data, lat_data, c=plot_c,vmin=plot_vmin, vmax=plot_vmax, cmap=plot_cmap, - s=plot_s,marker='o', edgecolors='none') - cnt, bins, art_hist = bk_plot.ax_hist.hist(alt_data, orientation='horizontal', - density=True, bins=80, range=(0, 20), color='black') - art_txt = plt.text(0.25, 0.10, str(len(alt_data)) + ' src', - fontsize='small', horizontalalignment='left', - verticalalignment='center',transform=bk_plot.ax_hist.transAxes) - art_out = [art_plan, art_th, art_lon, art_lat, art_txt] - # art_hist is a tuple of patch objects. Make it a flat list of artists - art_out.extend(art_hist) + s=plot_s, marker=marker, linewidths=edge_width, edgecolors=edge_color, **kwargs) + art_out = [art_plan, art_th, art_lon, art_lat] + + if add_to_histogram: + cnt, bins, art_hist = bk_plot.ax_hist.hist(alt_data, orientation='horizontal', + density=True, bins=80, range=(0, 20), color='black') + art_txt = plt.text(0.25, 0.10, str(len(alt_data)) + ' src', + fontsize='small', horizontalalignment='left', + verticalalignment='center',transform=bk_plot.ax_hist.transAxes) + # art_hist is a tuple of patch objects. Make it a flat list of artists + art_out.append(art_txt) + art_out.append(art_hist) + return art_out + + +def plot_2d_network_points(bk_plot, netw_data, actual_height=None, fake_ic_height=18, fake_cg_height=1, + color_by='time', pos_color='red', neg_color='blue', **kwargs): + """ + Plot points from a 2D lightning mapping neworks (ie, NLDN, ENTLN, etc) + + Parameters + ---------- + bk_plot : `pyxlma.plot.xlma_base_plot.BlankPlot` + A BlankPlot object to plot the data on + netw_data : `pandas.DataFrame` or `xarray.Dataset` + data object with columns/variables 'longitude', 'latitude', 'type' (CG/IC), and 'datetime' + actual_height : `numpy.ndarray` or `pandas.Series` or `xarray.DataArray` + the hieghts of the events to be plotted (default None, fake_ic_height and fake_cg_height used) + fake_ic_height : float + the altitude to plot IC points (default 18 km) + fake_cg_height : float + the altitude to plot CG points (default 1 km) + color_by : ['time', 'polarity'] + Whether to color the points by time or polarity. Default 'time'. Ignored if **kwargs contains 'c'. + pos_color : str + color for positive points (default 'red') if color_by='polarity' + neg_color : str + color for negative points (default 'blue') if color_by='polarity' + **kwargs + additional keyword arguments to pass to plt.scatter + + Returns + ------- + art_out : list + nested lists of artists created by plot_points (first list CG, second list IC) + + """ + + plot_c = kwargs.pop('c', None) + vmin = kwargs.pop('vmin', None) + vmax = kwargs.pop('vmax', None) + marker = kwargs.pop('marker', '^') + if actual_height is not None: + netw_data['height'] = actual_height + + if plot_c is not None: + netw_data['plot_c'] = plot_c + elif color_by == 'time': + netw_data['plot_c'] = color_by_time(netw_data.datetime, bk_plot.tlim)[2] + elif color_by == 'polarity': + pass + else: + raise ValueError("color_by must be 'time' or 'polarity'") + + cgs = netw_data[netw_data['type']=='CG'].copy() + ics = netw_data[netw_data['type']=='IC'].copy() + + if actual_height is None: + cgs['height'] = np.full_like(cgs.longitude, fake_cg_height) + ics['height'] = np.full_like(ics.longitude, fake_ic_height) + art_out = [] + if color_by == 'polarity': + cgpos = cgs[cgs.peak_current_kA>0] + cgneg = cgs[cgs.peak_current_kA<0] + icpos = ics[ics.peak_current_kA>0] + icneg = ics[ics.peak_current_kA<0] + art_out.append(plot_points(bk_plot, cgneg.longitude, cgneg.latitude, cgneg.height, + cgneg.datetime, c=neg_color, marker=marker, add_to_histogram=False, **kwargs)) + art_out.append(plot_points(bk_plot, cgpos.longitude, cgpos.latitude, cgpos.height, + cgpos.datetime, c=pos_color, marker=marker, add_to_histogram=False, **kwargs)) + art_out.append(plot_points(bk_plot, icneg.longitude, icneg.latitude, icneg.height, + icneg.datetime, c=neg_color, marker=marker, add_to_histogram=False, **kwargs)) + art_out.append(plot_points(bk_plot, icpos.longitude, icpos.latitude, icpos.height, + icpos.datetime, c=pos_color, marker=marker, add_to_histogram=False, **kwargs)) + else: + art_out.append(plot_points(bk_plot, cgs.longitude, cgs.latitude, cgs.height, + cgs.datetime, c=cgs.plot_c, vmin=vmin, vmax=vmax, marker=marker, add_to_histogram=False, **kwargs)) + art_out.append(plot_points(bk_plot, ics.longitude, ics.latitude, ics.height, + ics.datetime, c=ics.plot_c, vmin=vmin, vmax=vmax, marker=marker, add_to_histogram=False, **kwargs)) + return art_out + + +def plot_glm_events(glm, bk_plot, fake_alt=[0, 1], should_parallax_correct=True, poly_kwargs={}, vlines_kwargs={}): + """ + Plot event-level data from a glmtools dataset on a pyxlma.plot.xlma_base_plot.BlankPlot object. + Events that occupy the same pixel have their energies summed and plotted on the planview axis, event locations + are plotted on the lat/lon/time axes with an altitude specified as fake_alt. + Requires glmtools to be installed. + + Parameters + ---------- + glm : `xarray.Dataset` + A glmtools glm dataset to plot + bk_plot : `pyxlma.plot.xlma_base_plot.BlankPlot` + A BlankPlot object to plot the data on + fake_alt : list + the axes relative coordinates to plot the vertical lines for GLM events in the cross section, default [0, 1], + the full height of the axes. + should_parallax_correct : bool + whether to correct the GLM event locations for parallax effect. See https://doi.org/10.1029/2019JD030874 for more information. + poly_kwargs : dict + dictionary of additional keyword arguments to be passed to matplotlib Polygon + vlines_kwargs : dict + dictionary of additional keyword arguments to be passed to matplotlib vlines + + Returns + ------- + art_out : list + Handle to matplotlib polygon collection for the planview axis + """ + from cartopy import crs as ccrs + from glmtools.io.glm import get_lutevents + from glmtools.io.ccd import load_pixel_corner_lookup, quads_from_corner_lookup + from glmtools.io.lightning_ellipse import lightning_ellipse_rev, ltg_ellpse_rev + from pyxlma.coords import GeostationaryFixedGridSystem, GeographicSystem + + unique_ds = get_lutevents(glm) + evrad = unique_ds.lutevent_energy + + x_lut, y_lut, corner_lut = load_pixel_corner_lookup() + x_lut = x_lut * 1.0e-6 + y_lut = y_lut * 1.0e-6 + corner_lut = corner_lut*1e-6 + + event_polys = quads_from_corner_lookup(x_lut, y_lut, corner_lut, + unique_ds.lutevent_x, unique_ds.lutevent_y) + + glm['lutevent_corner_x'] = xr.DataArray(event_polys[:,:,0], dims=['lutevent_id', 'number_of_corners']) + glm['lutevent_corner_y'] = xr.DataArray(event_polys[:,:,1], dims=['lutevent_id', 'number_of_corners']) + + cx = glm.lutevent_corner_x + cy = glm.lutevent_corner_y + cz = np.zeros_like(cx) + sat_ecef_height = glm.nominal_satellite_height.data.astype(np.float64)*1000 + if should_parallax_correct: + ellps_rev_ver = ltg_ellpse_rev(glm.product_time.data.astype('datetime64[s]').item()) + ltg_ellps_re, ltg_ellps_rp = lightning_ellipse_rev[ellps_rev_ver] + gfgs_ellipse = [ltg_ellps_re, ltg_ellps_rp] + else: + gfgs_ellipse = 'WGS84' + ltg_ellps_re = None + ltg_ellps_rp = None + geofixcs = GeostationaryFixedGridSystem(subsat_lon=glm.lon_field_of_view.data.item(), ellipse=gfgs_ellipse, + sweep_axis='x', sat_ecef_height=sat_ecef_height) + grs80lla = GeographicSystem() + ltg_lon, ltg_lat, ltg_alt = grs80lla.fromECEF(*geofixcs.toECEF(cx,cy,cz)) + poly_verts = [] + for polynum in range(ltg_lon.shape[0]): + poly_lons = ltg_lon[polynum, :] + poly_lats = ltg_lat[polynum, :] + this_poly_verts = np.vstack([poly_lons, poly_lats]).T + poly_verts.append(this_poly_verts) + if hasattr(bk_plot.ax_plan, 'projection'): + map_proj = bk_plot.ax_plan.projection + if map_proj == ccrs.PlateCarree(): + transformed_pv = poly_verts + else: + transformed_pv = [map_proj.transform_points(ccrs.PlateCarree(), + this_poly_verts[:, 0], + this_poly_verts[:, 1])[:, 0:2] + for this_poly_verts in poly_verts] + else: + transformed_pv = poly_verts + patches = [Polygon(pv, closed=True) for pv in transformed_pv] + pc = PatchCollection(patches, **poly_kwargs) + pc.set_array(evrad.data) + bk_plot.ax_plan.add_collection(pc) + th_handle = bk_plot.ax_th.vlines(glm.event_time_offset.data, fake_alt[0], fake_alt[1], transform=bk_plot.ax_th.get_xaxis_transform(), **vlines_kwargs) + lon_handle = bk_plot.ax_lon.vlines(glm.event_lon, fake_alt[0], fake_alt[1], transform=bk_plot.ax_lon.get_xaxis_transform(), **vlines_kwargs) + lat_handle = bk_plot.ax_lat.hlines(glm.event_lat, fake_alt[0], fake_alt[1], transform=bk_plot.ax_lat.get_yaxis_transform(), **vlines_kwargs) + art_out = [pc, th_handle, lon_handle, lat_handle] return art_out + def plot_3d_grid(bk_plot, xedges, yedges, zedges, tedges, alt_lon, alt_lat, alt_time, lat_lon, - alt_data, plot_cmap): + alt_data, plot_cmap=None, **kwargs): """ Plot gridded fields on an existing bk_plot given x,y,z,t grids and respective grid edges + + In previous versions, 'plot_cmap' was required positional argument, this now defaults to None/matplotlib default unless overridden + Before the addition of **kwargs, 'vmin' was hardcoded to 0. This allows the user to specify a vmin in **kwargs, but maintain + backwards compatibility with assuming a vmin of 0 if no vmin is provided + """ + + plot_cmap = kwargs.pop('cmap', plot_cmap) + plot_vmin = kwargs.pop('vmin', 0) + alt_lon[alt_lon==0]=np.nan alt_lat[alt_lat==0]=np.nan lat_lon[lat_lon==0]=np.nan alt_time[alt_time==0]=np.nan - bk_plot.ax_lon.pcolormesh( xedges, zedges, alt_lon.T, cmap=plot_cmap, vmin=0) - bk_plot.ax_lat.pcolormesh( zedges, yedges, alt_lat.T, cmap=plot_cmap, vmin=0) - bk_plot.ax_plan.pcolormesh(xedges, yedges, lat_lon.T, cmap=plot_cmap, vmin=0) - bk_plot.ax_th.pcolormesh( tedges, zedges, alt_time.T, cmap=plot_cmap, vmin=0) + bk_plot.ax_lon.pcolormesh( xedges, zedges, alt_lon.T, cmap=plot_cmap, vmin=plot_vmin, **kwargs) + bk_plot.ax_lat.pcolormesh( zedges, yedges, alt_lat.T, cmap=plot_cmap, vmin=plot_vmin, **kwargs) + bk_plot.ax_plan.pcolormesh(xedges, yedges, lat_lon.T, cmap=plot_cmap, vmin=plot_vmin, **kwargs) + bk_plot.ax_th.pcolormesh( tedges, zedges, alt_time.T, cmap=plot_cmap, vmin=plot_vmin, **kwargs) bk_plot.ax_hist.hist(alt_data, orientation='horizontal', density=True, bins=80, range=(0, 20)) plt.text(0.25, 0.10, str(len(alt_data)) + ' src', diff --git a/pyxlma/xarray_util.py b/pyxlma/xarray_util.py new file mode 100644 index 0000000..79b67cc --- /dev/null +++ b/pyxlma/xarray_util.py @@ -0,0 +1,126 @@ +from collections import defaultdict +import xarray as xr +import numpy as np + +def get_1d_dims(d): + """ + Find all dimensions in a dataset that are purely 1-dimensional, + i.e., those dimensions that are not part of a 2D or higher-D + variable. + + arguments + d: xarray Dataset + returns + dims1d: a list of dimension names + """ + # Assume all dims coorespond to 1D vars + dims1d = list(d.dims.keys()) + for varname, var in d.variables.items(): + if len(var.dims) > 1: + for vardim in var.dims: + if vardim in dims1d: + dims1d.remove(str(vardim)) + return dims1d + +def gen_1d_datasets(d): + """ + Generate a sequence of datasets having only those variables + along each dimension that is only used for 1-dimensional variables. + + arguments + d: xarray Dataset + returns + generator function yielding a sequence of single-dimension datasets + """ + dims1d = get_1d_dims(d) +# print(dims1d) + for dim in dims1d: + all_dims = list(d.dims.keys()) + all_dims.remove(dim) + yield d.drop_dims(all_dims) + +def get_1d_datasets(d): + """ + Generate a list of datasets having only those variables + along each dimension that is only used for 1-dimensional variables. + + arguments + d: xarray Dataset + returns + a list of single-dimension datasets + """ + return [d1 for d1 in gen_1d_datasets(d, *args, **kwargs)] + +def get_scalar_vars(d): + scalars = [] + for varname, var in d.variables.items(): + if len(var.dims) == 0: + scalars.append(varname) + return scalars + +def concat_1d_dims(datasets, stack_scalars=None): + """ + For each xarray Dataset in datasets, concatenate (preserving the order of datasets) + all variables along dimensions that are only used for one-dimensional variables. + + arguments + d: iterable of xarray Datasets + stack_scalars: create a new dimension named with this value + that aggregates all scalar variables and coordinates + returns + a new xarray Dataset with only the single-dimension variables + """ + # dictionary mapping dimension names to a list of all + # datasets having only that dimension + all_1d_datasets = defaultdict(list) + + for d in datasets: + scalars = get_scalar_vars(d) + for d_1d_initial in gen_1d_datasets(d): + # Get rid of scalars + d_1d = d_1d_initial.drop(scalars) + dims = tuple(d_1d.dims.keys()) + all_1d_datasets[dims[0]].append(d_1d) + if stack_scalars: + # restore scalars along new dimension stack_scalars + scalar_dataset = xr.Dataset() + for scalar_var in scalars: + # promote from scalar to an array with a dimension, and remove + # the coordinate info so that it's just a regular variable. + as_1d = d[scalar_var].expand_dims(stack_scalars).reset_coords(drop=True) + scalar_dataset[scalar_var] = as_1d # xr.DataArray(as_1d, dims=[stack_scalars]) + all_1d_datasets[stack_scalars].append(scalar_dataset) + + unified = xr.Dataset() + for dim in all_1d_datasets: + combined = xr.concat(all_1d_datasets[dim], dim, coords='minimal', data_vars='minimal') + unified.update(combined) + return unified + +# datasets=[] +# for i, size in enumerate((4, 6)): +# a = xr.DataArray(10*i + np.arange(size), dims='x') +# b = xr.DataArray(10*i + np.arange(size/2), dims='y') +# c = xr.DataArray(20*i + np.arange(size*3), dims='t') +# d = xr.DataArray(11*i + np.arange(size*3), dims='t') +# T = xr.DataArray(10*i + np.arange(size)**2, dims='x') +# D = xr.DataArray(10*i + np.arange(size/2)**2, dims='y') +# z = xr.DataArray(10*i + np.arange(size*4)**2, dims='z') +# u = xr.DataArray(10*i + np.arange(size*5)**2, dims='u') +# v = xr.DataArray(12*i + np.arange(size*5)**2, dims='u') +# P = xr.DataArray(10*i + np.ones((size,int(size/2))), dims=['x', 'y']) +# Q = xr.DataArray(20*i + np.ones((size,int(size/2))), dims=['x', 'y']) +# d = xr.Dataset({'x':a,'y':b, 't':c, 'd':d, 'u':u, 'v':v, 'z':z, 'T':T, 'D':D, 'P':P, 'Q':Q}) +# datasets.append(d) +# # datasets.append(d[{'x':slice(None, None), 'y':slice(0,0)}]) +# for d in datasets: print(d,'\n') +# # xr.combine_by_coords(datasets, coords='all') +# # xr.combine_nested(datasets, coords='all', data_vars='all') + +# # print(get_1d_dims(d)) +# assert(get_1d_dims(d)==['t', 'u', 'z']) +# # for d1 in get_1d_datasets(d): +# # print(d1,'\n') + +# combined = concat_1d_dims(datasets) +# print(combined) \ No newline at end of file diff --git a/tests/test-env.yml b/tests/test-env.yml new file mode 100644 index 0000000..6a4aeaa --- /dev/null +++ b/tests/test-env.yml @@ -0,0 +1,19 @@ +name: pyxlma-tests +channels: +- conda-forge +dependencies: +- pytest-cov +- pytest-mpl +- xarray +- netcdf4 +- pandas +- numpy +- scipy +- scikit-learn +- pyproj +- metpy +- ipywidgets +- pip: + - git+https://github.com/deeplycloudy/lmatools + - git+https://github.com/deeplycloudy/glmtools@energy_threshold_util + - -e ../ \ No newline at end of file diff --git a/tests/test_base_plot.py b/tests/test_base_plot.py new file mode 100644 index 0000000..d350c75 --- /dev/null +++ b/tests/test_base_plot.py @@ -0,0 +1,208 @@ +import pytest +import xarray as xr +from pyxlma.plot.xlma_base_plot import * +from pyxlma.plot.xlma_plot_feature import * +from pyxlma.lmalib.grid import * +from pyxlma.lmalib.io import read as lma_read +from glmtools.io.glm import GLMDataset +import datetime as dt +import pandas as pd +import matplotlib.dates as md + +@pytest.mark.mpl_image_compare +def test_blank_plot(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + return bk_plot.fig + +@pytest.mark.mpl_image_compare +def test_blank_plot_labeled(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + subplot_labels(bk_plot) + return bk_plot.fig + +@pytest.mark.mpl_image_compare +def test_plot_feature_plot_points_positional(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + times = pd.Series(dataset.event_time.data.flatten()) + vmin, vmax, colors = color_by_time(pd.to_datetime(times), (start_time, end_time)) + plot_points(bk_plot, dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data/1000, + dataset.event_time.data, 'rainbow', 5, vmin, vmax, colors, 'k', 0.1, True) + return bk_plot.fig + +@pytest.mark.mpl_image_compare +def test_plot_feature_plot_points_old_kw(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + times = pd.Series(dataset.event_time.data.flatten()) + vmin, vmax, colors = color_by_time(pd.to_datetime(times), (start_time, end_time)) + plot_points(bk_plot, dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data/1000, + dataset.event_time.data, plot_cmap='rainbow', plot_s=5, plot_vmin=vmin, plot_vmax=vmax, + plot_c=colors, edge_color='k', edge_width=0.1, add_to_histogram=True) + return bk_plot.fig + +@pytest.mark.mpl_image_compare +def test_plot_feature_plot_points_new_kw(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + times = pd.Series(dataset.event_time.data.flatten()) + vmin, vmax, colors = color_by_time(pd.to_datetime(times), (start_time, end_time)) + plot_points(bk_plot, dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data/1000, + dataset.event_time.data, cmap='rainbow', s=5, vmin=vmin, vmax=vmax, + c=colors, edgecolors='k', linewidths=0.1, add_to_histogram=True) + return bk_plot.fig + + +@pytest.mark.mpl_image_compare +def test_plot_feature_plot_points_new_kw_no_bkmap(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=False, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + times = pd.Series(dataset.event_time.data.flatten()) + vmin, vmax, colors = color_by_time(pd.to_datetime(times), (start_time, end_time)) + plot_points(bk_plot, dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data/1000, + dataset.event_time.data, cmap='rainbow', s=5, vmin=vmin, vmax=vmax, + c=colors, edgecolors='k', linewidths=0.1, add_to_histogram=True) + return bk_plot.fig + + +@pytest.mark.mpl_image_compare +def test_plot_feature_plot_points_new_kw_no_hist(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + times = pd.Series(dataset.event_time.data.flatten()) + vmin, vmax, colors = color_by_time(pd.to_datetime(times), (start_time, end_time)) + plot_points(bk_plot, dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data/1000, + dataset.event_time.data, cmap='rainbow', s=5, vmin=vmin, vmax=vmax, + c=colors, edgecolors='k', linewidths=0.1, add_to_histogram=False) + return bk_plot.fig + + +@pytest.mark.mpl_image_compare +def test_plot_feature_plot_3d_grid(): + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + x_edges = np.linspace(-103.5, -99.5, 100) + y_edges = np.linspace(31.5, 35.5, 100) + z_edges = np.linspace(0, 20, 100) + t_edges = md.date2num(pd.date_range(start='2023-12-24T00:57:00', end='2023-12-24T00:58:00', periods=7).values) + histograms = setup_hist(dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data/1000, + dataset.event_time.data, x_edges, y_edges, z_edges, t_edges) + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + + plot_3d_grid(bk_plot, x_edges, y_edges, z_edges, t_edges, *histograms, dataset.event_altitude.data/1000, cmap='plasma') + return bk_plot.fig + +@pytest.mark.mpl_image_compare +def test_plot_feature_plot_3d_grid_old_kw(): + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + x_edges = np.linspace(-103.5, -99.5, 100) + y_edges = np.linspace(31.5, 35.5, 100) + z_edges = np.linspace(0, 20, 100) + t_edges = md.date2num(pd.date_range(start='2023-12-24T00:57:00', end='2023-12-24T00:58:00', periods=7).values) + histograms = setup_hist(dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data/1000, + dataset.event_time.data, x_edges, y_edges, z_edges, t_edges) + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + + plot_3d_grid(bk_plot, x_edges, y_edges, z_edges, t_edges, *histograms, dataset.event_altitude.data/1000, plot_cmap='plasma') + return bk_plot.fig + + +@pytest.mark.mpl_image_compare +def test_plot_feature_inset_view(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + times = pd.Series(dataset.event_time.data.flatten()) + vmin, vmax, colors = color_by_time(pd.to_datetime(times), (start_time, end_time)) + plot_points(bk_plot, dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data/1000, + dataset.event_time.data, cmap='rainbow', s=5, vmin=vmin, vmax=vmax, + c=colors, edgecolors='k', linewidths=0.1, add_to_histogram=True) + inset_view(bk_plot, dataset.event_longitude.data, dataset.event_latitude.data, + [-102.75, -102.25], [32, 32.5], .01, .01) + return bk_plot.fig + + +@pytest.mark.mpl_image_compare +def test_plot_feature_ndln_time(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + nldn_data = lma_read.nldn('examples/network_samples/gld360enldnns_20231224_daily_v1_lit.raw') + plot_2d_network_points(bk_plot, nldn_data) + return bk_plot.fig + + +def test_plot_feature_ndln_bad_arg(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + nldn_data = lma_read.nldn('examples/network_samples/gld360enldnns_20231224_daily_v1_lit.raw') + with pytest.raises(ValueError, match="color_by must be 'time' or 'polarity'"): + plot_2d_network_points(bk_plot, nldn_data, color_by='bad_arg') + + +@pytest.mark.mpl_image_compare +def test_plot_feature_ndln_custom_colors(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + nldn_data = lma_read.nldn('examples/network_samples/gld360enldnns_20231224_daily_v1_lit.raw') + plot_2d_network_points(bk_plot, nldn_data, c=[0, 10, 5], cmap='plasma', vmin=0, vmax=10) + return bk_plot.fig + + +@pytest.mark.mpl_image_compare +def test_plot_feature_entln_real_height(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + entln_data = lma_read.entln('examples/network_samples/lxarchive_pulse20231224.csv') + plot_2d_network_points(bk_plot, entln_data, actual_height=entln_data['icheight']) + return bk_plot.fig + + +@pytest.mark.mpl_image_compare +def test_plot_feature_ndln_polarity(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + nldn_data = lma_read.nldn('examples/network_samples/gld360enldnns_20231224_daily_v1_lit.raw') + plot_2d_network_points(bk_plot, nldn_data, color_by='polarity') + return bk_plot.fig + + +@pytest.mark.mpl_image_compare +def test_plot_feature_glm_events(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + glm_data = GLMDataset('examples/network_samples/OR_GLM-L2-LCFA_G16_s20233580057000_e20233580057200_c20233580057222.nc') + plot_glm_events(glm_data.dataset, bk_plot, fake_alt=[0, 0.25], poly_kwargs={'cmap' : 'plasma'}, vlines_kwargs={'linewidths' : 0.5}, should_parallax_correct=False) + return bk_plot.fig + +@pytest.mark.mpl_image_compare +def test_plot_feature_glm_events_parallax(): + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + end_time = start_time + dt.timedelta(seconds=60) + bk_plot = BlankPlot(start_time, bkgmap=True, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], tlim=[start_time, end_time], title='XLMA Test Plot') + glm_data = GLMDataset('examples/network_samples/OR_GLM-L2-LCFA_G16_s20233580057000_e20233580057200_c20233580057222.nc') + plot_glm_events(glm_data.dataset, bk_plot, fake_alt=[0, 0.25], poly_kwargs={'cmap' : 'plasma'}, vlines_kwargs={'linewidths' : 0.5}, should_parallax_correct=True) + return bk_plot.fig diff --git a/tests/test_coords.py b/tests/test_coords.py new file mode 100644 index 0000000..f84f321 --- /dev/null +++ b/tests/test_coords.py @@ -0,0 +1,149 @@ +from pyxlma.coords import * +import pytest +import numpy as np +from sklearn.neighbors import KDTree + +test_lats = np.array([33.5, 1.0, 0.0, 0.0, 0.0, 10.0, -10.0, 33.606968]) +test_lons = np.array([-101.5, -75.0, -85.0, -65.0, -75.0, -75.0, -75.0, -101.822625]) +test_alts = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 984.0]) + +test_ecef_X = np.array([-1061448.75418035, 1650533.58831094, 555891.26758132, + 2695517.17208404, 1650783.32787306, 1625868.32721344, + 1625868.32721344, -1089633.44245767]) +test_ecef_Y = np.array([-5217187.30723133, -6159875.21117539, -6353866.26310279, + -5780555.22988658, -6160807.25190988, -6067823.20357756, + -6067823.20357756, -5205511.43302535]) +test_ecef_Z = np.array([3500334.28802236, 110568.77482457, 0, 0, + 0, 1100248.54773536, -1100248.54773536, 3510766.26631805]) + +def test_geographic(): + geosys = GeographicSystem() + ecef_coords = geosys.toECEF(test_lons, test_lats, test_alts) + lons, lats, alts = geosys.fromECEF(*ecef_coords) + assert np.allclose(ecef_coords[0], test_ecef_X) + assert np.allclose(ecef_coords[1], test_ecef_Y) + assert np.allclose(ecef_coords[2], test_ecef_Z) + assert np.allclose(lons, test_lons) + assert np.allclose(lats, test_lats) + assert np.allclose(alts, test_alts) + +def test_geographic_one_point(): + geosys = GeographicSystem() + ecef_coords = geosys.toECEF(np.atleast_1d(test_lons[-1]), np.atleast_1d(test_lats[-1]), np.atleast_1d(test_alts[-1])) + lons, lats, alts = geosys.fromECEF(*ecef_coords) + assert np.allclose(ecef_coords[0], test_ecef_X[-1]) + assert np.allclose(ecef_coords[1], test_ecef_Y[-1]) + assert np.allclose(ecef_coords[2], test_ecef_Z[-1]) + assert np.allclose(lons[0], test_lons[-1]) + assert np.allclose(lats[0], test_lats[-1]) + assert np.allclose(alts[0], test_alts[-1]) + +def test_geographic_custom_r_both(): + geosys = GeographicSystem(r_equator=6378.137, r_pole=6356.752) + ecef_coords = geosys.toECEF(test_lons, test_lats, test_alts) + lons, lats, alts = geosys.fromECEF(*ecef_coords) + assert np.allclose(lons, test_lons) + assert np.allclose(lats, test_lats) + assert np.allclose(alts, test_alts) + +def test_geographic_custom_r_eq(): + geosys = GeographicSystem(r_equator=6378.137) + ecef_coords = geosys.toECEF(test_lons, test_lats, test_alts) + lons, lats, alts = geosys.fromECEF(*ecef_coords) + assert np.allclose(lons, test_lons) + assert np.allclose(lats, test_lats) + assert np.allclose(alts, test_alts) + +def test_geographic_custom_r_pole(): + geosys = GeographicSystem(r_pole=6356.752) + ecef_coords = geosys.toECEF(test_lons, test_lats, test_alts) + lons, lats, alts = geosys.fromECEF(*ecef_coords) + assert np.allclose(lons, test_lons) + assert np.allclose(lats, test_lats) + assert np.allclose(alts, test_alts) + +def test_equidistant_cylindrical(): + eqsys = MapProjection() + ecef_coords = eqsys.toECEF(0, 0, 0) + x, y, z = eqsys.fromECEF(*ecef_coords) + assert np.allclose(x, 0) + assert np.allclose(y, 0) + assert np.allclose(z, 0) + +def test_equidistant_cylindrical_custom_point(): + eqsys = MapProjection(ctrLat=test_lats[-1], ctrLon=test_lons[-1]) + ecef_coords = eqsys.toECEF(0, 0, 0) + x, y, z = eqsys.fromECEF(*ecef_coords) + assert np.allclose(x, 0) + assert np.allclose(y, 0) + assert np.allclose(z, 0) + + +# def test_px_grid(): +# lon = np.arange(-105, -99.9, 0.5) +# x_coord = np.arange(0, len(lon)) +# lat = np.arange(30, 35.1, 0.5) +# y_coord = np.arange(0, len(lat)) +# lon, lat = np.meshgrid(lon, lat) +# pxgrid = PixelGrid(lon, lat, KDTree, x_coord, y_coord) +# ecef_coords = pxgrid.toECEF(np.array(7), np.array(7), np.array(0)) +# x, y, z = pxgrid.fromECEF(*ecef_coords) + +# assert np.allclose(ecef_coords[0], test_ecef_X[0]) +# assert np.allclose(ecef_coords[1], test_ecef_Y[0]) +# assert np.allclose(ecef_coords[2], test_ecef_Z[0]) + +# assert np.allclose(x, 7) +# assert np.allclose(y, 7) +# assert np.allclose(z, 0) + +def test_satellite_system(): + sat = GeostationaryFixedGridSystem(subsat_lon=75.2) + ecef_coords = sat.toECEF(0.01, 0.01, 0.01) + x, y, z = sat.fromECEF(*ecef_coords) + + assert np.allclose(x, 0.01) + assert np.allclose(y, 0.01) + assert np.allclose(z, 0.01) + + +def test_radar_system_height(): + ADRAD_rcs = RadarCoordinateSystem(30.6177, -96.3365, 114) + tornado_ground_range, beam_height_agl = ADRAD_rcs.getGroundRangeHeight(17150, 1.4) + assert np.allclose(tornado_ground_range, 17144.013390611748) + assert np.allclose(beam_height_agl, 550.2784673999995) + +def test_radar_system_elevation(): + ADRAD_rcs = RadarCoordinateSystem(30.6177, -96.3365, 114) + tornado_slant_range, radar_elevation = ADRAD_rcs.getSlantRangeElevation(17144.013390611748, 550.2784673999995) + assert np.allclose(tornado_slant_range, 17150) + assert np.allclose(radar_elevation, 1.4) + +def test_radar_system_lla(): + ADRAD_rcs = RadarCoordinateSystem(30.6177, -96.3365, 114) + tornado_lon, tornado_lat, tornado_alt = ADRAD_rcs.toLonLatAlt(np.atleast_1d(17150), np.atleast_1d(228), np.atleast_1d(1.4)) + assert np.allclose(tornado_lat, 30.51415605367721) + assert np.allclose(tornado_lon, -96.46923405085701) + assert np.allclose(tornado_alt, 550.2784674) + +def test_radar_system_ecef(): + ADRAD_rcs = RadarCoordinateSystem(30.6177, -96.3365, 114) + tornado_x, tornado_y, tornado_z = ADRAD_rcs.toECEF(np.atleast_1d(17150), np.atleast_1d(228), np.atleast_1d(1.4)) + tornado_r, tornado_az, tornado_el = ADRAD_rcs.fromECEF(tornado_x, tornado_y, tornado_z) + assert np.allclose(tornado_r, 17150) + assert np.allclose(tornado_az, 228) + assert np.allclose(tornado_el, 1.4) + +def test_tpcs(): + tpcs = TangentPlaneCartesianSystem(ctrLat=test_lats[-1], ctrLon=test_lons[-1], ctrAlt=test_alts[-1]) + ecef_coords = tpcs.toECEF(100, 100, 100) + x, y, z = tpcs.fromECEF(*ecef_coords) + assert np.allclose(x, 100) + assert np.allclose(y, 100) + assert np.allclose(z, 100) + +def test_tpcs_local(): + tpcs = TangentPlaneCartesianSystem(ctrLat=test_lats[-1], ctrLon=test_lons[-1], ctrAlt=test_alts[-1]) + ecef_coords = tpcs.toLocal(np.array([[100, 100, 100], [200, 200, 200], [300, 300, 300]])) + local_coords = tpcs.fromLocal(ecef_coords) + assert np.allclose(local_coords, np.array([[100, 100, 100], [200, 200, 200], [300, 300, 300]])) diff --git a/tests/test_flash.py b/tests/test_flash.py new file mode 100644 index 0000000..c3a4597 --- /dev/null +++ b/tests/test_flash.py @@ -0,0 +1,79 @@ +"""Test functionality of pyxlma.lmalib.flash""" + +import xarray as xr +import numpy as np +import pytest +from pyxlma.lmalib.flash.cluster import cluster_flashes +from pyxlma.lmalib.flash.properties import * + + +def compare_dataarrays(tocheck, truth, var, rtol=1.e-5, atol=1.e-8): + """Compare two dataarrays""" + if truth[var].data.dtype == 'datetime64[ns]' or truth[var].data.dtype == 'timedelta64[ns]': + if tocheck[var].data.dtype == 'float64': + truth[var].data = truth[var].data.astype(float)/1e9 + np.testing.assert_allclose(tocheck[var].data.astype(float), truth[var].data.astype(float), rtol=rtol, atol=atol, equal_nan=True) + else: + np.testing.assert_allclose(tocheck[var].data, truth[var].data, rtol=rtol, atol=atol, equal_nan=True) + + +def test_cluster_flashes(): + """Test clustering of flashes""" + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + clustered = cluster_flashes(dataset) + truth = xr.open_dataset('tests/truth/lma_netcdf/lma_clustered.nc') + for var in truth.data_vars: + compare_dataarrays(clustered, truth, var) + + +def test_flash_stats(): + """Test calculation of flash statistics""" + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + stats = flash_stats(cluster_flashes(dataset)) + truth = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc') + for var in truth.data_vars: + compare_dataarrays(stats, truth, var) + + +def test_filter_flashes(): + """Test filtering of flashes""" + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc') + filtered = filter_flashes(dataset, flash_event_count=(100, 500)) + assert np.min(filtered.flash_event_count.data) >= 100 + assert np.max(filtered.flash_event_count.data) <= 500 + + truth = xr.open_dataset('tests/truth/lma_netcdf/lma_filtered.nc') + + for var in truth.data_vars: + compare_dataarrays(filtered, truth, var) + + +def test_flilter_flashes_no_prune(): + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc') + filtered = filter_flashes(dataset, flash_event_count=(100, 500), prune=False) + assert np.all(filtered.event_id.data == dataset.event_id.data) + assert np.min(filtered.flash_event_count.data) >= 100 + assert np.max(filtered.flash_event_count.data) <= 500 + + +def test_filter_no_stats(): + """"Test filtering of flashes without flash statistics""" + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + dataset = cluster_flashes(dataset) + with pytest.raises(ValueError, match='Before filtering a dataset by flash properties, call flash_stats on the dataset to compute flash properties.'): + filter_flashes(dataset, flash_event_count=(100, 500)) + + +def test_event_area(): + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + x, y, z = local_cartesian(dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data, + dataset.network_center_longitude.data, dataset.network_center_latitude.data, dataset.network_center_altitude.data) + assert np.isclose(event_hull_area(x, y, z), 5491450433206.501) + + +def test_event_volume(): + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + x, y, z = local_cartesian(dataset.event_longitude.data, dataset.event_latitude.data, dataset.event_altitude.data, + dataset.network_center_longitude.data, dataset.network_center_latitude.data, dataset.network_center_altitude.data) + assert np.isclose(event_hull_volume(x[0:10], y[0:10], z[0:10]), 56753729942.624825) + \ No newline at end of file diff --git a/tests/test_grid.py b/tests/test_grid.py new file mode 100644 index 0000000..9845a75 --- /dev/null +++ b/tests/test_grid.py @@ -0,0 +1,69 @@ +import xarray as xr +from pyxlma.lmalib.grid import * +from tests.test_flash import compare_dataarrays + + +def test_create_regular_grid(): + """Test creation of regular grid""" + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc') + grid_range = 0.5 + grid_h_res = 0.1 + grid_height = 20 + grid_v_res = 1 + lon_range = (dataset.network_center_longitude.data.item() - grid_range, dataset.network_center_longitude.data.item() + grid_range, grid_h_res) + lat_range = (dataset.network_center_latitude.data.item() - grid_range, dataset.network_center_latitude.data.item() + grid_range, grid_h_res) + alt_range = (0, grid_height, grid_v_res) + time_range = (dataset.event_time.data.min(), dataset.event_time.data.max(), np.timedelta64(1, 'm')) + grid_edge_ranges ={ + 'grid_longitude_edge':lon_range, + 'grid_latitude_edge':lat_range, + 'grid_altitude_edge':alt_range, + 'grid_time_edge':time_range, + } + grid_center_names ={ + 'grid_longitude_edge':'grid_longitude', + 'grid_latitude_edge':'grid_latitude', + 'grid_altitude_edge':'grid_altitude', + 'grid_time_edge':'grid_time', + } + empty_grid = create_regular_grid(grid_edge_ranges, grid_center_names) + xr.testing.assert_equal(empty_grid, xr.open_dataset('tests/truth/lma_netcdf/empty_grid.nc')) + +def test_assign_regular_bins(): + """Test assigning lightning data to regular bins""" + empty_grid = xr.open_dataset('tests/truth/lma_netcdf/empty_grid.nc') + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc') + event_coord_names = { + 'event_longitude':'grid_longitude_edge', + 'event_latitude':'grid_latitude_edge', + 'event_altitude':'grid_altitude_edge', + 'event_time':'grid_time_edge', + } + binned_grid = assign_regular_bins(empty_grid, dataset, event_coord_names, append_indices=True) + truth = xr.open_dataset('tests/truth/lma_netcdf/binned_grid.nc') + + for var in truth.data_vars: + compare_dataarrays(binned_grid, truth, var) + + +def test_events_to_grid(): + """Test gridding lightning data""" + empty_grid = xr.open_dataset('tests/truth/lma_netcdf/empty_grid.nc') + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc') + event_coord_names = { + 'event_longitude':'grid_longitude_edge', + 'event_latitude':'grid_latitude_edge', + 'event_altitude':'grid_altitude_edge', + 'event_time':'grid_time_edge', + } + binned_grid = assign_regular_bins(empty_grid, dataset, event_coord_names, append_indices=True) + + grid_spatial_coords=('grid_time', 'grid_altitude', 'grid_latitude', 'grid_longitude') + event_spatial_vars = ('event_time', 'event_altitude', 'event_latitude', 'event_longitude') + + gridded_lma = events_to_grid(binned_grid, empty_grid, pixel_id_var='pixel_id', event_spatial_vars=event_spatial_vars, grid_spatial_coords=grid_spatial_coords) + + truth = xr.open_dataset('tests/truth/lma_netcdf/gridded_lma.nc') + + for var in truth.data_vars: + compare_dataarrays(gridded_lma, truth, var) diff --git a/tests/test_interactive.py b/tests/test_interactive.py new file mode 100644 index 0000000..cc88eb9 --- /dev/null +++ b/tests/test_interactive.py @@ -0,0 +1,14 @@ +import pytest +import xarray as xr +import datetime as dt +from pyxlma.plot.interactive import * + +@pytest.mark.mpl_image_compare +def test_interactive(): + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + buffer = 2.5 + xlim = (dataset.network_center_longitude.data-buffer, dataset.network_center_longitude.data+buffer) + ylim = (dataset.network_center_latitude.data-buffer, dataset.network_center_latitude.data+buffer) + zlim = (0, 20) + interact = InteractiveLMAPlot(dataset, clon=dataset.network_center_longitude, clat=dataset.network_center_latitude, xlim=xlim, ylim=ylim, zlim=zlim, plot_cmap='rainbow') + return interact.lma_plot.fig diff --git a/tests/test_intercept_rhi.py b/tests/test_intercept_rhi.py new file mode 100644 index 0000000..fa1ceaf --- /dev/null +++ b/tests/test_intercept_rhi.py @@ -0,0 +1,69 @@ +from pyxlma.lmalib import lma_intercept_rhi +from pyxlma.lmalib.io import read as lma_read +from datetime import timedelta +from os import listdir +import xarray as xr +import numpy as np + + +def test_intercept_rhi(): + """Test functionality of pyxlma.plot.lma_intercept_rhi""" + files_to_read = listdir('examples/data/') + files_to_read = ['examples/data/'+file for file in files_to_read] + dataset, start_date = lma_read.dataset(files_to_read) + lma_radar_range, lma_plane_distance, lma_arl, event_mask = lma_intercept_rhi.find_lma_points_near_rhi(dataset, 33.56, -101.81, 1600.0, 225, start_date+timedelta(seconds=30), distance_threshold=500, time_threshold=15) + lma_indices = np.flatnonzero(event_mask) + true_radar_range = np.array([-24314.95714263, -24255.42134556, -23947.00309296, -24201.07999466, + -24216.22494555, -24209.0779143 , -24252.54770026, -24196.90100367, + -24261.77984816, -24002.07311165, -23977.02349918, -23922.33835871, + -24061.79138268, -24063.22463418, -24095.4865483 , -23629.07408446, + -23711.71702149, -23815.42007446, -23801.22424582, -23811.93829348, + -23880.1053328 , -23787.41575292, -23714.05514112, -23753.40801102, + -23697.50367439, -23808.79529204, -23702.79936852, -23733.26491339, + -23654.19223296, -23639.08173578, -23684.40998865, -22516.2312817 , + -24956.31925763, -25087.31684706, -25278.91196761, -25284.01320286, + -25287.53041298, -25404.67473694, -25398.64205462, -25328.81651605, + -25326.06791898, -25236.99678994, -25662.94030104, -25703.02511737, + -25733.48865903, -25548.97709685, -25840.7709326 , -26074.465209 , + -26135.94920909, -25718.73533927, -26070.13443815, -25742.36675196, + -26088.72774852, -26150.28911721, -22628.68886763]) + true_plane_distance = np.array([476.82409578, 471.16428085, 324.56896919, 464.99685801, + 443.51557595, 420.92058199, 399.02641209, 300.20108514, + 330.2928808 , 284.48486544, 277.49727128, 228.24504655, + 101.86785646, 46.09758876, 159.49307647, 422.34413364, + 462.82523707, 479.04669236, 450.06522663, 495.67782222, + 458.62470004, 397.71230587, 322.04159886, 358.54206568, + 263.69453322, 215.27970084, 234.5031954 , 274.60201124, + 184.85941046, 153.78011609, 203.25043091, 141.91307406, + 291.47165215, 274.21197041, 378.27769326, 296.90514221, + 280.02141178, 354.41770628, 327.55116845, 385.75338378, + 177.16633758, 35.34312855, 287.23074474, 220.35513643, + 134.09140466, 13.21239026, 82.51552308, 6.20537471, + 46.62220121, 275.7219309 , 67.74536475, 254.6569667 , + 175.35763511, 190.40871913, 100.09991281]) + true_radar_arl = np.array([2427.10820635, 2622.80703538, 2642.20087632, 2740.2745904 , + 2930.63971071, 2752.31680521, 2927.29418501, 2859.50970951, + 3625.92759368, 2942.6473481 , 2931.8120999 , 2970.00909331, + 3111.63730722, 3062.02183736, 2859.75475096, 1394.13388557, + 3537.06953818, 3000.85787932, 2971.55268989, 3552.81355203, + 2836.86606787, 3226.92886031, 2460.49052586, 3200.61752493, + 2790.06678895, 2551.85152503, 2572.98642808, 3500.51837957, + 3063.97138001, 3031.9081397 , 3487.58191853, 6265.1033319 , + 2473.16902922, 2203.98364055, 2856.44738566, 2345.22649489, + 2515.22474203, 2635.02689176, 2923.88449116, 2959.58956895, + 2263.4423628 , 2731.91960772, 2772.38586661, 2878.88741945, + 2746.43494287, 2584.35605103, 2341.75645168, 2613.90859694, + 2919.30959546, 2614.69497102, 2473.25507194, 2782.82223582, + 2461.79799166, 2377.74468257, 7332.69246886]) + true_lma_ids = np.array([12731, 12733, 12734, 12736, 12741, 12743, 12745, 12752, 12756, + 12768, 12769, 12775, 12787, 12788, 12797, 12991, 13002, 13005, + 13006, 13007, 13008, 13009, 13010, 13011, 13012, 13013, 13015, + 13016, 13023, 13029, 13030, 13201, 13388, 13390, 13395, 13396, + 13397, 13399, 13402, 13405, 13408, 13409, 13411, 13412, 13415, + 13419, 13428, 13437, 13439, 13442, 13443, 13444, 13448, 13451, + 13556], dtype='uint64') + assert np.allclose(lma_radar_range, true_radar_range, atol=1e-3) + assert np.allclose(lma_plane_distance, true_plane_distance, atol=1e-3) + assert np.allclose(lma_arl, true_radar_arl, atol=1e-3) + assert np.all(lma_indices == true_lma_ids) + diff --git a/tests/test_io.py b/tests/test_io.py new file mode 100644 index 0000000..85e8d74 --- /dev/null +++ b/tests/test_io.py @@ -0,0 +1,62 @@ +"""Test functionality of pyxlma.lmalib.io""" + +from os import listdir +from datetime import datetime as dt +import xarray as xr +import pandas as pd +import numpy as np +from pyxlma.lmalib.io import read as lma_read +from tests.test_flash import compare_dataarrays + + +def test_read_mf_dataset(): + """Test reading of multiple .dat.gz files into a single lma dataset""" + files_to_read = listdir('examples/data/') + files_to_read = ['examples/data/'+file for file in files_to_read] + dataset, start_date = lma_read.dataset(files_to_read) + assert start_date == dt(2023, 12, 24, 0, 57, 1) + assert dataset == xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + + +def test_read_onefile_dataset(): + dataset, start_date = lma_read.dataset('examples/data/WTLMA_231224_005701_0001.dat.gz') + assert start_date == dt(2023, 12, 24, 0, 57, 1) + truth = xr.open_dataset('tests/truth/lma_netcdf/small_lma.nc') + for var in truth.data_vars: + compare_dataarrays(dataset, truth, var) + + +def test_read_nldn(): + dataset = lma_read.nldn('examples/network_samples/gld360enldnns_20231224_daily_v1_lit.raw') + truth = pd.DataFrame({ + 'latitude' : [33.582, 33.590, 33.585], + 'longitude' : [-101.881, -102.032, -101.872], + 'peak_current_kA' : [12.0, -9.0, -15.0], + 'multiplicity' : [0, 0, 0], + 'semimajor': [0.4, 0.2, 2.2], + 'semiminor': [0.2, 0.2, 1.5], + 'majorminorratio': [1.8, 1.0, 1.5], + 'ellipseangle' : [75, 3, 26], + 'chi2' : [0.8, 1.0, 2.1], + 'num_stations': [4, 5, 6], + 'type': ['IC', 'CG', 'IC'], + 'datetime' : [np.datetime64('2023-12-24T00:57:01.123456789'), np.datetime64('2023-12-24T00:57:31.987654321'), np.datetime64('2023-12-24T00:57:58.135792468')] + }) + assert dataset.equals(truth) + + +def test_read_entln(): + dataset = lma_read.entln('examples/network_samples/lxarchive_pulse20231224.csv') + truth = pd.DataFrame({ + 'type': ['IC', 'CG', 'IC'], + 'datetime' : [np.datetime64('2023-12-24T00:57:04.123456789'), np.datetime64('2023-12-24T00:57:26.987654321'), np.datetime64('2023-12-24T00:57:47.246813579')], + 'latitude' : [33.581914, 33.590077, 33.584480], + 'longitude' : [-101.880986, -102.032033, -101.871498], + 'peak_current_kA' : [12.345, 9.876, 15.79], + 'icheight' : [3014, 6028, 13591], + 'num_stations': [8, 7, 11], + 'ellipseangle' : [104., 99., 102.], + 'semimajor': [0.1855, 0.3465, 0.089], + 'semiminor': [0.029, 0.054, 0.031] + }) + assert dataset.equals(truth) diff --git a/tests/test_plot_feature.py b/tests/test_plot_feature.py new file mode 100644 index 0000000..9823f7f --- /dev/null +++ b/tests/test_plot_feature.py @@ -0,0 +1,52 @@ +from pyxlma.plot.xlma_plot_feature import * +import xarray as xr +import numpy as np +from datetime import datetime as dt + +def test_subset(): + lma = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + lon_subset, lat_subset, alt_subset, time_subset, selection = subset(lma.event_longitude.data, lma.event_latitude.data, lma.event_altitude.data, + lma.event_time.data, lma.event_chi2.data, lma.event_stations.data, + (-101.7, -101.4), (33.4, 34.8), (0, 3000), + (np.datetime64('2023-12-24T00:57:00'), np.datetime64('2023-12-24T00:57:10')), 1, 8) + assert np.allclose(lon_subset, np.array([-101.59106, -101.598236, -101.59939, -101.59875, -101.601425, -101.60555, -101.60554, -101.60838, -101.60368, -101.62052])) + assert np.allclose(lat_subset, np.array([33.68953, 33.684334, 33.68242, 33.67924, 33.678104, 33.676983, 33.675335, 33.677456, 33.67102, 33.666958])) + assert np.allclose(alt_subset, np.array([2974.67, 2986.99, 2936.03, 2920.17, 2797.01, 2933.09, 2659.97, 2886.72, 2716.22, 2943.72])) + assert np.allclose(time_subset.astype(float), np.array(['2023-12-24T00:57:07.731986515', '2023-12-24T00:57:07.800978747', '2023-12-24T00:57:07.803362858', '2023-12-24T00:57:07.805963963', + '2023-12-24T00:57:07.806720943', '2023-12-24T00:57:07.809493631', '2023-12-24T00:57:07.810448100', '2023-12-24T00:57:07.811465266', + '2023-12-24T00:57:07.814960674', '2023-12-24T00:57:07.826344209']).astype(np.datetime64).astype(float)) + assert np.sum(selection) == 10 + +def test_color_by_time_datetime_nolimit(): + some_datetimes = np.array([dt(2021, 4, 9, 1, 51, 0), dt(2021, 4, 9, 1, 52, 0), dt(2021, 4, 9, 1, 53, 0), dt(2021, 4, 9, 1, 54, 0), dt(2021, 4, 9, 1, 59, 0)]) + vmin, vmax, colors = color_by_time(some_datetimes) + assert vmin == 0 + assert vmax == 480 + assert np.allclose(colors, np.array([0, 60, 120, 180, 480])) + + +def test_color_by_time_datetime_limit(): + some_datetimes = np.array([dt(2021, 4, 9, 1, 51, 0), dt(2021, 4, 9, 1, 52, 0), dt(2021, 4, 9, 1, 53, 0), dt(2021, 4, 9, 1, 54, 0), dt(2021, 4, 9, 1, 59, 0)]) + limits = [dt(2021, 4, 9, 1, 50, 0), dt(2021, 4, 9, 2, 0, 0)] + vmin, vmax, colors = color_by_time(some_datetimes, limits) + assert vmin == 0 + assert vmax == 600 + assert np.allclose(colors, np.array([60, 120, 180, 240, 540])) + +def test_color_by_time_xarray(): + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + vmin, vmax, colors = color_by_time(dataset.event_time) + assert vmin == 0 + assert np.isclose(vmax, 57.943683385849) + assert np.isclose(np.mean(colors), 30.483982899376258) + assert np.isclose(np.std(colors), 17.25687093241869) + +def test_setup_hist(): + lma = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + alt_lon, alt_lat, alt_time, lat_lon = setup_hist(lma.event_longitude.data, lma.event_latitude.data, lma.event_altitude.data, + lma.event_time.data, 2, 2, 2, 2) + + assert np.allclose(alt_lon, np.array([[21082, 1], [0, 1]])) + assert np.allclose(alt_lat, np.array([[5, 21077], [1, 1]])) + assert np.allclose(alt_time, np.array([[9991, 1], [11091, 1]])) + assert np.allclose(lat_lon, np.array([[6, 21077], [0, 1]])) diff --git a/tests/test_xlma_plot.py b/tests/test_xlma_plot.py new file mode 100644 index 0000000..83ce099 --- /dev/null +++ b/tests/test_xlma_plot.py @@ -0,0 +1,13 @@ +import pytest +from pyxlma.plot.xlma import * +import xarray as xr +import datetime as dt +import pandas as pd + +@pytest.mark.mpl_image_compare +def test_xlma_plot(): + dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc') + start_time = dt.datetime(2023, 12, 24, 0, 57, 0, 0) + print(pd.Series(dataset.event_time.data.flatten())) + xlma = XlmaPlot(dataset, start_time, True, True, 'xarray', False, xlim=[-103.5, -99.5], ylim=[31.5, 35.5], zlim=[0, 20], cmap='rainbow', chi2=1) + return xlma.fig \ No newline at end of file diff --git a/tests/truth/images/test_blank_plot.png b/tests/truth/images/test_blank_plot.png new file mode 100644 index 0000000..40070e2 Binary files /dev/null and b/tests/truth/images/test_blank_plot.png differ diff --git a/tests/truth/images/test_blank_plot_labeled.png b/tests/truth/images/test_blank_plot_labeled.png new file mode 100644 index 0000000..765b029 Binary files /dev/null and b/tests/truth/images/test_blank_plot_labeled.png differ diff --git a/tests/truth/images/test_interactive.png b/tests/truth/images/test_interactive.png new file mode 100644 index 0000000..6bae648 Binary files /dev/null and b/tests/truth/images/test_interactive.png differ diff --git a/tests/truth/images/test_plot_feature_entln_real_height.png b/tests/truth/images/test_plot_feature_entln_real_height.png new file mode 100644 index 0000000..0ebf657 Binary files /dev/null and b/tests/truth/images/test_plot_feature_entln_real_height.png differ diff --git a/tests/truth/images/test_plot_feature_glm_events.png b/tests/truth/images/test_plot_feature_glm_events.png new file mode 100644 index 0000000..05c57ff Binary files /dev/null and b/tests/truth/images/test_plot_feature_glm_events.png differ diff --git a/tests/truth/images/test_plot_feature_glm_events_parallax.png b/tests/truth/images/test_plot_feature_glm_events_parallax.png new file mode 100644 index 0000000..6d75313 Binary files /dev/null and b/tests/truth/images/test_plot_feature_glm_events_parallax.png differ diff --git a/tests/truth/images/test_plot_feature_inset_view.png b/tests/truth/images/test_plot_feature_inset_view.png new file mode 100644 index 0000000..df33fb6 Binary files /dev/null and b/tests/truth/images/test_plot_feature_inset_view.png differ diff --git a/tests/truth/images/test_plot_feature_ndln_custom_colors.png b/tests/truth/images/test_plot_feature_ndln_custom_colors.png new file mode 100644 index 0000000..c7e80b7 Binary files /dev/null and b/tests/truth/images/test_plot_feature_ndln_custom_colors.png differ diff --git a/tests/truth/images/test_plot_feature_ndln_polarity.png b/tests/truth/images/test_plot_feature_ndln_polarity.png new file mode 100644 index 0000000..28537a4 Binary files /dev/null and b/tests/truth/images/test_plot_feature_ndln_polarity.png differ diff --git a/tests/truth/images/test_plot_feature_ndln_time.png b/tests/truth/images/test_plot_feature_ndln_time.png new file mode 100644 index 0000000..2cf6c86 Binary files /dev/null and b/tests/truth/images/test_plot_feature_ndln_time.png differ diff --git a/tests/truth/images/test_plot_feature_plot_3d_grid.png b/tests/truth/images/test_plot_feature_plot_3d_grid.png new file mode 100644 index 0000000..8934c29 Binary files /dev/null and b/tests/truth/images/test_plot_feature_plot_3d_grid.png differ diff --git a/tests/truth/images/test_plot_feature_plot_3d_grid_old_kw.png b/tests/truth/images/test_plot_feature_plot_3d_grid_old_kw.png new file mode 100644 index 0000000..8934c29 Binary files /dev/null and b/tests/truth/images/test_plot_feature_plot_3d_grid_old_kw.png differ diff --git a/tests/truth/images/test_plot_feature_plot_points_new_kw.png b/tests/truth/images/test_plot_feature_plot_points_new_kw.png new file mode 100644 index 0000000..9602645 Binary files /dev/null and b/tests/truth/images/test_plot_feature_plot_points_new_kw.png differ diff --git a/tests/truth/images/test_plot_feature_plot_points_new_kw_no_bkmap.png b/tests/truth/images/test_plot_feature_plot_points_new_kw_no_bkmap.png new file mode 100644 index 0000000..824e106 Binary files /dev/null and b/tests/truth/images/test_plot_feature_plot_points_new_kw_no_bkmap.png differ diff --git a/tests/truth/images/test_plot_feature_plot_points_new_kw_no_hist.png b/tests/truth/images/test_plot_feature_plot_points_new_kw_no_hist.png new file mode 100644 index 0000000..bcbe0d7 Binary files /dev/null and b/tests/truth/images/test_plot_feature_plot_points_new_kw_no_hist.png differ diff --git a/tests/truth/images/test_plot_feature_plot_points_old_kw.png b/tests/truth/images/test_plot_feature_plot_points_old_kw.png new file mode 100644 index 0000000..9602645 Binary files /dev/null and b/tests/truth/images/test_plot_feature_plot_points_old_kw.png differ diff --git a/tests/truth/images/test_plot_feature_plot_points_positional.png b/tests/truth/images/test_plot_feature_plot_points_positional.png new file mode 100644 index 0000000..9602645 Binary files /dev/null and b/tests/truth/images/test_plot_feature_plot_points_positional.png differ diff --git a/tests/truth/images/test_xlma_plot.png b/tests/truth/images/test_xlma_plot.png new file mode 100644 index 0000000..4e3a8d6 Binary files /dev/null and b/tests/truth/images/test_xlma_plot.png differ diff --git a/tests/truth/lma_netcdf/binned_grid.nc b/tests/truth/lma_netcdf/binned_grid.nc new file mode 100644 index 0000000..172df29 Binary files /dev/null and b/tests/truth/lma_netcdf/binned_grid.nc differ diff --git a/tests/truth/lma_netcdf/empty_grid.nc b/tests/truth/lma_netcdf/empty_grid.nc new file mode 100644 index 0000000..21e2b41 Binary files /dev/null and b/tests/truth/lma_netcdf/empty_grid.nc differ diff --git a/tests/truth/lma_netcdf/gridded_lma.nc b/tests/truth/lma_netcdf/gridded_lma.nc new file mode 100644 index 0000000..181afde Binary files /dev/null and b/tests/truth/lma_netcdf/gridded_lma.nc differ diff --git a/tests/truth/lma_netcdf/lma.nc b/tests/truth/lma_netcdf/lma.nc new file mode 100644 index 0000000..998088e Binary files /dev/null and b/tests/truth/lma_netcdf/lma.nc differ diff --git a/tests/truth/lma_netcdf/lma_clustered.nc b/tests/truth/lma_netcdf/lma_clustered.nc new file mode 100644 index 0000000..d543dcb Binary files /dev/null and b/tests/truth/lma_netcdf/lma_clustered.nc differ diff --git a/tests/truth/lma_netcdf/lma_filtered.nc b/tests/truth/lma_netcdf/lma_filtered.nc new file mode 100644 index 0000000..dcf9598 Binary files /dev/null and b/tests/truth/lma_netcdf/lma_filtered.nc differ diff --git a/tests/truth/lma_netcdf/lma_stats.nc b/tests/truth/lma_netcdf/lma_stats.nc new file mode 100644 index 0000000..24b4072 Binary files /dev/null and b/tests/truth/lma_netcdf/lma_stats.nc differ diff --git a/tests/truth/lma_netcdf/small_lma.nc b/tests/truth/lma_netcdf/small_lma.nc new file mode 100644 index 0000000..21a1872 Binary files /dev/null and b/tests/truth/lma_netcdf/small_lma.nc differ