diff --git a/input/era5_smoothed_topography_land_masks/era-spectral_T10_16x32.nc b/input/era5_smoothed_topography_land_masks/era-spectral_T10_16x32.nc new file mode 100644 index 000000000..e3dcea8f8 Binary files /dev/null and b/input/era5_smoothed_topography_land_masks/era-spectral_T10_16x32.nc differ diff --git a/input/era5_smoothed_topography_land_masks/era-spectral_T170_256x512.nc b/input/era5_smoothed_topography_land_masks/era-spectral_T170_256x512.nc new file mode 100644 index 000000000..74f66e17c Binary files /dev/null and b/input/era5_smoothed_topography_land_masks/era-spectral_T170_256x512.nc differ diff --git a/input/era5_smoothed_topography_land_masks/era-spectral_T21_32x64.nc b/input/era5_smoothed_topography_land_masks/era-spectral_T21_32x64.nc new file mode 100644 index 000000000..4891c276c Binary files /dev/null and b/input/era5_smoothed_topography_land_masks/era-spectral_T21_32x64.nc differ diff --git a/input/era5_smoothed_topography_land_masks/era-spectral_T341_512x1024.nc b/input/era5_smoothed_topography_land_masks/era-spectral_T341_512x1024.nc new file mode 100644 index 000000000..1914cdef4 Binary files /dev/null and b/input/era5_smoothed_topography_land_masks/era-spectral_T341_512x1024.nc differ diff --git a/input/era5_smoothed_topography_land_masks/era-spectral_T42_64x128.nc b/input/era5_smoothed_topography_land_masks/era-spectral_T42_64x128.nc new file mode 100644 index 000000000..55ca0643a Binary files /dev/null and b/input/era5_smoothed_topography_land_masks/era-spectral_T42_64x128.nc differ diff --git a/input/era5_smoothed_topography_land_masks/era-spectral_T85_128x256.nc b/input/era5_smoothed_topography_land_masks/era-spectral_T85_128x256.nc new file mode 100644 index 000000000..0d4ec65ab Binary files /dev/null and b/input/era5_smoothed_topography_land_masks/era-spectral_T85_128x256.nc differ diff --git a/src/extra/python/scripts/create_era5_topography.py b/src/extra/python/scripts/create_era5_topography.py new file mode 100644 index 000000000..006e96e2d --- /dev/null +++ b/src/extra/python/scripts/create_era5_topography.py @@ -0,0 +1,189 @@ +import numpy as np +import xarray as xr +import pyshtools as pysh +from numpy.polynomial.legendre import leggauss +from scipy.interpolate import RegularGridInterpolator +import copy, os +from scipy.special import j1 +from resolutions import get_grid_for_truncation +from rich.console import Console +from rich.table import Table + +# Physical constants +g = 9.80 # m/s^2, Earth gravity +a = 6376.0e3 # m, Earth radius + +# Global console object for Rich output +console = Console() + +def load_era5_invariant_fields(file_dir: str, file_lsm: str, file_z: str) -> xr.Dataset: + """Load and merge ERA5 invariant fields (land-sea mask and topography).""" + lsm_path = os.path.join(file_dir, file_lsm) + z_path = os.path.join(file_dir, file_z) + + if not os.path.exists(lsm_path) or not os.path.exists(z_path): + console.print("[bold red]Error:[/bold red] ERA5 invariant files not found.") + console.print(f"Expected paths:\n- {lsm_path}\n- {z_path}") + console.print("\n[bold]Note:[/bold] You must download the ERA5 invariant fields manually.") + console.print("Please refer to the README.md for download instructions.\n") + raise FileNotFoundError("ERA5 invariant files not found.") + + era5_lsm = xr.open_dataset(file_dir + file_lsm).squeeze() + era5_z = xr.open_dataset(file_dir + file_z).squeeze() + + if era5_lsm.lsm.shape != era5_z.z.shape: + raise ValueError( + f"Shape mismatch between geopotential (z: {era5_z.z.shape}) " + f"and land-sea mask (lsm: {era5_lsm.lsm.shape})" + ) + + invar = xr.merge([era5_lsm, era5_z], compat='no_conflicts') + invar = invar.assign(zsurf=invar.z / g) + invar = invar.drop_vars(['z']) + + return invar + +def print_era5_statistics(ds: xr.Dataset) -> None: + """Print statistics for ERA5 invariant fields.""" + nlat = len(ds['latitude']) + nlon = len(ds['longitude']) + + console.print("[bold]ERA5 Invariant Fields Statistics[/bold]") + console.print(f"[bold]Grid Shape (lat, lon):[/bold] {nlat} x {nlon}\n") + + table = Table(title="Variable Statistics", show_header=True, header_style="bold") + table.add_column("Variable", style="cyan", no_wrap=True) + table.add_column("Minimum") + table.add_column("Maximum") + table.add_row("zsurf", f"{ds.zsurf.min():.1f}", f"{ds.zsurf.max():.1f}") + table.add_row("lsm", f"{ds.lsm.min():.1f}", f"{ds.lsm.max():.1f}") + + console.print(table) + +def gaussian_grid(n_lat: int, n_lon: int) -> tuple[np.ndarray, np.ndarray]: + """ + Compute a Gaussian latitude-longitude grid. + Gaussian latitudes are the roots of the Legendre polynomial of degree n_lat, + converted to degrees. Longitudes are equally spaced. + """ + gauss_points, _ = leggauss(n_lat) + lats = (np.arcsin(gauss_points) * (180.0 / np.pi))[::-1] + lons = np.linspace(0, 360, n_lon, endpoint=False) + return lats, lons + +def dh_sh_filter( + ds: xr.Dataset, + tnum: int, + n_target_lat: int = None, + n_target_lon: int = None, + gaussian: bool = False, +) -> xr.Dataset: + """ + Perform spherical harmonic filtering on an input dataset and return fields + on a specified latitude-longitude grid. + """ + if tnum is None: + raise ValueError("`tnum` must be provided as a positive, non-zero integer.") + if not isinstance(tnum, int) or tnum <= 0: + raise ValueError("`tnum` must be a positive, non-zero integer.") + if (n_target_lat is None) != (n_target_lon is None): + raise ValueError("`n_target_lat` and `n_target_lon` must be provided as a pair or not at all.") + + # If tnum is provided but n_target_lat and n_target_lon are not, use get_grid_for_truncation + if tnum is not None and n_target_lat is None and n_target_lon is None: + grid_params = get_grid_for_truncation(tnum) + n_target_lat = grid_params['nlat'] + n_target_lon = grid_params['nlon'] + console.print(f"[bold]Note:[/bold] Using grid parameters for T{tnum}: nlat={n_target_lat}, nlon={n_target_lon}") + + nlat_orig = len(ds.latitude) + nlon_orig = len(ds.longitude) + + # pyshtools expects input grids to match certain formats, one of which is + # a DH grid, which has shape (x, x/2) rather than the supplied (x, x/2 + 1) + lats_dh = np.linspace(90, -90, (nlat_orig // 2) * 2) + lons_dh = np.linspace(0, 360, nlon_orig, endpoint=False) + + # Define target lat-lon arrays (can be uniform or Gaussian) + if gaussian: + target_lats, target_lons = gaussian_grid(n_target_lat, n_target_lon) + else: + target_lats = np.linspace(90, -90, n_target_lat) + target_lons = np.linspace(0, 360, n_target_lon, endpoint=False) + + # Ensure latitude ascending for xarray interp + ds_sorted = ds.sortby('latitude') + filtered_dict = {} + + for var in ds_sorted.data_vars: + da = ds_sorted[var] + # Interpolate input to DH grid (roughly matching input resolution) + da_interp = da.interp( + latitude=xr.DataArray(lats_dh, dims='latitude'), + longitude=xr.DataArray(lons_dh, dims='longitude'), + method='linear' + ) + # Create SHGrid object and expand coefficients upto tnum + data_dh = da_interp.values + grid = pysh.SHGrid.from_array(data_dh) + clm = grid.expand(lmax_calc=tnum) + # Apply filter + clm_filtered = copy.deepcopy(clm) + Theta_opt = 3.8317 / (tnum + 0.5) + for l in range(tnum + 1): + for m in range(l + 1): + factor = 2 * j1(l * Theta_opt) / (l * Theta_opt) if l > 0 else 1.0 + clm_filtered.coeffs[0, l, m] *= factor + clm_filtered.coeffs[1, l, m] *= factor + # Expand back to DH grid (PySh decides shape based on tnum) + filtered_grid = clm_filtered.expand() + filtered_array = filtered_grid.to_array() + lat_filtered = np.linspace(90, -90, filtered_array.shape[0]) + lon_filtered = np.linspace(0, 360, filtered_array.shape[1], endpoint=False) + # Wrap as xarray DataArray and interpolate to target grid + da_filtered = xr.DataArray(filtered_array, + dims=('latitude', 'longitude'), + coords={'latitude': lat_filtered, + 'longitude': lon_filtered}) + da_filtered_interp = da_filtered.interp(latitude=target_lats, + longitude=target_lons, + method='linear') + # Round values in land mask + if var == 'lsm': + da_filtered_interp = np.rint(da_filtered_interp) + filtered_dict[var] = (('latitude', 'longitude'), da_filtered_interp.values) + + coords_dict = {'latitude': target_lats, 'longitude': target_lons} + filtered_ds = xr.Dataset(filtered_dict, coords=coords_dict) + # Sort latitude ascending and rename + filtered_ds = filtered_ds.sortby('latitude', ascending=True) + filtered_ds = filtered_ds.astype({v: "float32" for v in filtered_ds.data_vars}) + filtered_ds = filtered_ds.rename({'latitude':'lat', 'longitude':'lon', 'lsm':'land_mask'}) + + return filtered_ds + +def generate_spectral_grids(ds: xr.Dataset, truncations: list[int]) -> None: + """Generate and save spectral grids for specified truncations.""" + for t in truncations: + gstat = get_grid_for_truncation(t) + grid_name = f'era-spectral_T{gstat["nfou"]}_{gstat["nlat"]}x{gstat["nlon"]}' + ds_tgrid = dh_sh_filter(ds, tnum=gstat["nfou"], gaussian=True) + ds_tgrid.to_netcdf(f"{grid_name}.nc") + console.print(f"Saved new T{gstat['nfou']} topography file [bold]{grid_name}.nc[/bold]\n") + +def main(): + # File paths + file_dir = 'era5_invariant_fields/' + file_lsm = 'ecmwf-era5_oper_an_sfc_200001010000.lsm.inv.nc' + file_z = 'ecmwf-era5_oper_an_sfc_200001010000.z.inv.nc' + + # Load and print statistics + invar = load_era5_invariant_fields(file_dir, file_lsm, file_z) + print_era5_statistics(invar) + + # Generate spectral grids + truncations = [10, 21, 42, 85, 170, 341] + generate_spectral_grids(invar, truncations) + +if __name__ == "__main__": + main() diff --git a/src/extra/python/scripts/era5_invariant_fields/README.md b/src/extra/python/scripts/era5_invariant_fields/README.md new file mode 100644 index 000000000..2b9304d54 --- /dev/null +++ b/src/extra/python/scripts/era5_invariant_fields/README.md @@ -0,0 +1,12 @@ +# Topography Source Files +## ERA-5 + +Isca provides a number of preconstructed topography files for use with common resolutions, which can be found under `Isca/input/era5_smoothed_topography_land_masks`. These were constructed with the script `Isca/src/extra/python/scripts/create_era5_topography.py` and make use of the ERA-5 invariant fields as an input. The ERA-5 invariant fields are not included as part of Isca, however they may be obtained by using either directly from the [CEDA](https://catalogue.ceda.ac.uk/uuid/2c8f38fac04945b89cf12d6e9c928c6f/) archive (recommended) or from the ERA-5 [CDS](https://cds.climate.copernicus.eu/datasets). Note that a login is required for both websites. + +If downloading from CEDA the following files are required: +- Topography: `ecmwf-era5_oper_an_sfc_200001010000.z.inv` +- Land-sea mask: `ecmwf-era5_oper_an_sfc_200001010000.lsm.inv` + +## Other Sources/Planets +In theory the script `create_era5_topography.py` may be used to create smoothed/filtered topography input files from other data sources. You may need to change input variable names or disable the land-sea mask, but in principle the topography of other planets (e.g. Mars) can be used. + diff --git a/src/extra/python/scripts/resolutions.py b/src/extra/python/scripts/resolutions.py old mode 100755 new mode 100644 index b636bf9cc..3d7dc6300 --- a/src/extra/python/scripts/resolutions.py +++ b/src/extra/python/scripts/resolutions.py @@ -1,49 +1,188 @@ -#!/usr/bin/python - import math +from typing import List, Optional +from rich.table import Table +from rich.console import Console +from rich.prompt import Prompt -def prime_factors(n): # based on http://stackoverflow.com/questions/15347174/python-finding-prime-factors +def prime_factors(n: int) -> List[int]: + """Compute the prime factors of an integer n.""" + factors: List[int] = [] i = 2 - factors = [] - while i*i <= n: - if n % i: - i += 1 - else: - n //= i + while i * i <= n: + while n % i == 0: factors.append(i) - if n > 1 or len(factors) == 0: + n //= i + i += 1 + if n > 1 or not factors: factors.append(n) return factors -lat_mult = 4 # choose from factors of nproc (e.g. 1 to find smallest nlat, nproc to load-balance) and/or factors matching lon (anti-aliasing: 2*nlat ~ nlon)? -lon_maxprime = 2 # small for efficient FFTs: no higher than 3 or 5? -latlon_maxprod = 1024*1024 # maximum grid size - -nlat = lat_mult -nlon = 1 -nfou = 1 - -print("Each row indicates, for the corresponding triangular truncation, the minimum lat & lon grid sizes") -print("satisfying both the anti-aliasing constraints and any (lat) multiple or (lon) factorization constraints.") -print("The T number is the minimum num_fourier but is set by num_spherical - 1 (e.g. num_spherical = 43 for T42).") -print("The percentage of the grid in excess of that required to satisfy only the anti-aliasing constraints is also shown.") -print("") -print("In this run: lat is a multiple of "+str(lat_mult)+", lon cannot have any prime factor exceeding "+str(lon_maxprime)) -print("") -print("T\tfou\tsph\tlat\tlon\tover%\tlon factorization") -print("") - -while 1: - nsph = nfou + 1 - while 2*nlat < 3*(nsph - 1) + 1: +def next_valid_lat(nlat: int, nsph: int, lat_mult: int) -> int: + """Find the smallest multiple of nlat satisfying the anti-aliasing condition.""" + while 2 * nlat < 3 * (nsph - 1) + 1: nlat += lat_mult - while nlon < 3*nfou + 1 or prime_factors(nlon)[-1] > lon_maxprime: + return nlat + +def next_valid_lon(nlon: int, nfou: int, lon_maxprime: int) -> int: + """Find the smallest nlon meeting FFT factorization and anti-aliasing constraints.""" + while nlon < 3 * nfou + 1 or max(prime_factors(nlon)) > lon_maxprime: nlon += 1 - size = nlat*nlon - if size > latlon_maxprod: - print("Next grid size ("+str(size)+") exceeds specified maximum "+str(latlon_maxprod)) - break - size_over = size - int(math.ceil(float(3*(nsph - 1) + 1)/2)*(3*nfou + 1)) - strover = " - " if size_over == 0 else ("%4.1f" % (float(100*size_over)/size)) - print("T"+str(nfou)+"\t"+str(nfou)+"\t"+str(nsph)+"\t"+str(nlat)+"\t"+str(nlon)+"\t"+strover+"\t"+str(prime_factors(nlon))) - nfou += 1 + return nlon + + +def get_grid_for_truncation( + nfou: int, + lat_mult: int = 4, + lon_maxprime: int = 2, + ) -> dict: + """Return grid parameters for a specific truncation number.""" + nlat, nlon = lat_mult, 1 + nsph = nfou + 1 + nlat = next_valid_lat(nlat, nsph, lat_mult) + nlon = next_valid_lon(nlon, nfou, lon_maxprime) + size = nlat * nlon + min_size = math.ceil((3 * (nsph - 1) + 1) / 2) * (3 * nfou + 1) + over = size - min_size + percent_over = " - " if over == 0 else f"{100 * over / size:4.1f}" + truncation_summary = { + "nfou": nfou, + "nsph": nsph, + "nlat": nlat, + "nlon": nlon, + "over%": percent_over, + "lon_factors": prime_factors(nlon), + } + return truncation_summary + + +def generate_grids_table( + nfou_max: Optional[int] = None, + lat_mult: int = 4, + lon_maxprime: int = 2, + latlon_maxprod: int = 1024 * 1024, + ) -> Table: + + """Generate a rich.Table for spectral transform grids.""" + table = Table(title="Spectral Transform Grids") + table.add_column("T") + table.add_column("fou") + table.add_column("sph") + table.add_column("lat") + table.add_column("lon") + table.add_column("over%") + table.add_column("lon factorization") + + nlat, nlon, nfou = lat_mult, 1, 1 + console = Console() + + while True: + if nfou_max and nfou > nfou_max: + break + nsph = nfou + 1 + nlat = next_valid_lat(nlat, nsph, lat_mult) + nlon = next_valid_lon(nlon, nfou, lon_maxprime) + size = nlat * nlon + if size > latlon_maxprod: + exceeded = True + break + min_size = math.ceil((3 * (nsph - 1) + 1) / 2) * (3 * nfou + 1) + over = size - min_size + percent_over = " - " if over == 0 else f"{100 * over / size:4.1f}" + table.add_row( + f"[bold cyan]T{nfou}[/bold cyan]", + str(nfou), + str(nsph), + str(nlat), + str(nlon), + percent_over, + ", ".join(map(str, prime_factors(nlon))), + ) + nfou += 1 + return table, exceeded, size, latlon_maxprod + + + +def main(): + console = Console() + console.print( + "[bold]Summary of Spectral Grid Sizes in Isca[/bold]\n" + "This script will show you the minimum lat-lon grid shapes used in a\n" + "given triangular truncation that satisfies both anti-aliasing and\n" + "any (lat) multiple or (lon) factorisation constraints.\n\n" + + + "[bold]Spectral Grid Summary[/bold]\n" + "This table shows the spectral transform grids for some typical\n" + "resolutions (e.g. T42, T85), with columns defined by:\n" + "- [teal]T[/teal]: Truncation number\n" + "- [teal]fou[/teal]: Fourier wavenumber\n" + "- [teal]sph[/teal]: Spherical harmonic degree\n" + "- [teal]lat[/teal]: Number of latitude grid points\n" + "- [teal]lon[/teal]: Number of longitude grid points\n" + "- [teal]over%[/teal]: Percentage over minimum required grid size required\n" + " to only satisfy anti-aliasing constraint\n" + "- [teal]lon factorization[/teal]: Prime factors of longitude resolution.\n" + ) + + key_resolutions = [21, 42, 85, 170] + table = Table(title="Common spectral transform grids (see below for more)") + table.add_column("T") + table.add_column("fou") + table.add_column("sph") + table.add_column("lat") + table.add_column("lon") + table.add_column("over%") + table.add_column("lon factorization") + + for nfou in key_resolutions: + grid = get_grid_for_truncation(nfou) + table.add_row( + f"[bold cyan]T{nfou}[/bold cyan]", + str(grid["nfou"]), + str(grid["nsph"]), + str(grid["nlat"]), + str(grid["nlon"]), + grid["over%"], + ", ".join(map(str, grid["lon_factors"])), + ) + + console.print(table) + + console.print( + "\n[bold]Options:[/bold]\n" + "a) Display full table of resolutions for all truncation numbers (T1 to Tx)\n" + "b) Print grid parameters for a single specified truncation number (e.g. T21, T42)\n" + "q) [yellow]Quit[/yellow]" + ) + + choice = Prompt.ask( + "\nEnter your choice", + choices=["a", "b", "q"], + default="q", + ) + + if choice == "a": + table, exceeded, size, latlon_maxprod = generate_grids_table() + console.print( + "\n[bold]Full Spectral Transform Grids Table[/bold]\n" + "This table shows all available resolutions for spectral transform grids.\n" + ) + console.print(table) + if exceeded: + console.print(f"[bold green]Next grid size ({size}) exceeds specified maximum {latlon_maxprod}[/bold green]") + elif choice == "b": + nfou = int(Prompt.ask("\nEnter truncation number (e.g., 21, 42, 85, 170)")) + grid = get_grid_for_truncation(nfou) + console.print( + f"\n[bold]Grid parameters for T{nfou}:[/bold]\n" + f"- Truncation number (T): {grid['nfou']}\n" + f"- Spherical harmonic degree (sph): {grid['nsph']}\n" + f"- Number of latitude grid points (lat): {grid['nlat']}\n" + f"- Number of longitude grid points (lon): {grid['nlon']}\n" + f"- Percentage over minimum grid size (over%): {grid['over%']}\n" + f"- Longitude prime factors: {', '.join(map(str, grid['lon_factors']))}\n" + ) + +if __name__ == "__main__": + main() +