#!/usr/bin/env python3 """ Rescale an atmospheric tracer of a dynamics start file to a target inventory. The tracer field is multiplied by the single factor that brings its global inventory to the requested value, so the shape of the distribution is kept. The inventory is mass-weighted with the air mass of each cell ("masse"), which is what makes a target such as "10 pr-um of water vapour" meaningful. Usage: rescale_tracer.py [start.nc] --tracer h2o_vap --target 10 --unit pr-um [--output restart.nc | --inplace] [--force] [--saturation] [--exner milieu|hyb] [--startfi startfi.nc] [--reservoir h2o_ice] [--molar-mass 18.01528e-3] The result goes to restart.nc, and to restartfi.nc for --startfi, which is what a run calls the state it produces: rename them to start.nc and startfi.nc to start from them. Without --target, the current inventory is reported and nothing is written. Without --tracer, every tracer of the file is reported. If no file path is provided, the script will prompt for one with tab-completion. """ ## @file rescale_tracer.py # @author JB Clement # @date 24/08/2026 import os import sys import glob import shutil import readline import argparse import warnings import numpy as np from netCDF4 import Dataset # netCDF4 1.7 sets the shape of the array it is handed in place on every write, # which NumPy 2.5 deprecates. Nothing on this side can avoid it, and the message # would only bury the report. warnings.filterwarnings('ignore', category=DeprecationWarning, message='Setting the shape on a NumPy array') # Entries of LMDZ's "controle" vector worth reading, as documented by # "controle_descriptor" and fixed by dynetat0. Same layout in every start, # restart and diagfi file. CONTROL_NAME = 'controle' CONTROL_RADIUS = 4 CONTROL_GRAVITY = 6 CONTROL_CPP = 7 CONTROL_KAPPA = 8 CONTROL_PREFF = 17 # Universal gas constant [J.mol-1.K-1] R_UNIVERSAL = 8.314462618 # Density of liquid water [kg.m-3], to turn a water column into precipitable # microns: 1 kg.m-2 = 1 mm = 1000 pr-um RHO_WATER = 1000. # Molar masses [kg.mol-1] of the species a tracer name may start with. Only # needed for the "mol/mol" unit, and always overridable with --molar-mass. MOLAR_MASSES = { 'h2o': 18.01528e-3, 'co2': 44.0095e-3, 'co': 28.0101e-3, 'o2': 31.9988e-3, 'o3': 47.9982e-3, 'n2': 28.0134e-3, 'ar': 39.948e-3, 'h2': 2.01588e-3, 'ch4': 16.0425e-3, 'oh': 17.0073e-3, 'ho2': 33.0067e-3, 'h2o2': 34.0147e-3, 'no': 30.0061e-3, 'no2': 46.0055e-3, 'n': 14.0067e-3, 'o': 15.9994e-3, 'h': 1.00794e-3, 'he': 4.002602e-3, } # Tracers whose mixing ratio a --saturation run is not allowed to push above # saturation. Ice tracers are not condensable vapours and are left alone. SATURATING = ('h2o_vap',) # Units a target may be expressed in UNITS = ('kg', 'kg/m2', 'pr-um', 'kg/kg', 'mol/mol') # Variables carried on the tracer grid that are not tracers NOT_TRACERS = ('teta', 'masse', 'ucov', 'vcov', 'phi', 'pk') def fail(message): """ Print an error message in red and leave with a non-zero status. """ print(f"\033[91mError: {message}\033[0m") sys.exit(1) def warn(message): """ Print a warning in yellow and carry on. """ print(f"\033[93mCaution: {message}\033[0m") def complete_filename(text, state): """ Tab-completion function for readline: completes filesystem paths. Appends '/' if the match is a directory. """ # The text forms a partial path; glob for matching entries if "*" not in text: text_glob = text + "*" else: text_glob = text matches = glob.glob(os.path.expanduser(text_glob)) # Add a trailing slash for directories matches = [m + "/" if os.path.isdir(m) else m for m in matches] try: return matches[state] except IndexError: return None def read_array(ds, name): """ Read a variable as a plain float64 array, masked entries turned into NaNs. """ if name not in ds.variables: fail(f"variable \"{name}\" not found in \"{ds.filepath()}\"!") data = ds.variables[name][...] if np.ma.isMaskedArray(data): data = data.filled(np.nan) return np.asarray(data, dtype=np.float64) def control(ds, index, what): """ One entry of LMDZ's "controle" vector, or None when the file has no such vector or it is too short. 'what' only serves the error message. """ if CONTROL_NAME not in ds.variables: return None values = read_array(ds, CONTROL_NAME).ravel() if values.size <= index: return None value = float(values[index]) if not np.isfinite(value): fail(f"\"{CONTROL_NAME}\" holds a non-finite {what}!") return value class Grid: """ The geometry and the physical constants of a dynamics start file, and the air mass of every cell, which every inventory is weighted with. The lon-lat grid of the dynamics is redundant: the last longitude repeats the first one, and both pole rows repeat the same value over all longitudes with "aire" already holding one nlon-th of the polar cap. A global sum therefore runs over the first nlon columns and over nothing else. """ def __init__(self, ds): self.ds = ds self.path = ds.filepath() for name in ('aire', 'ps', 'longitude', 'latitude'): if name not in ds.variables: fail(f"\"{self.path}\" has no \"{name}\": this does not look " f"like a dynamics start file!") lon = read_array(ds, 'longitude') # Drop the duplicated last longitude, when the axis wraps around self.nlon = lon.size - 1 if abs(lon[-1] - lon[0] - 360.) < 1.e-6 else lon.size self.aire = read_array(ds, 'aire') self.nlat = self.aire.shape[0] self.area_total = float(self.aire[:, :self.nlon].sum()) self.radius = control(ds, CONTROL_RADIUS, 'radius') self.gravity = control(ds, CONTROL_GRAVITY, 'gravity') self.cpp = control(ds, CONTROL_CPP, 'specific heat') self.kappa = control(ds, CONTROL_KAPPA, 'kappa') self.preff = control(ds, CONTROL_PREFF, 'reference pressure') if self.gravity is None or self.gravity <= 0.: fail(f"\"{self.path}\" gives no usable gravity in \"{CONTROL_NAME}\"!") # Mean molar mass of the air, from the file's own constants: # kappa = r/cpp and r = R/M, so M = R/(kappa*cpp). A generic run on # another planet is thus converted with that planet's atmosphere. self.molar_mass_air = None if self.kappa and self.cpp: self.molar_mass_air = R_UNIVERSAL/(self.kappa*self.cpp) self.ps = read_array(ds, 'ps') self.ntime = self.ps.shape[0] self.masse = self.cell_mass() self.nlev = self.masse.shape[1] self.air_mass = self.total(self.masse) def cell_mass(self): """ The air mass of each cell [kg], shaped (Time, altitude, lat, lon). Taken from "masse" when the file carries it, and rebuilt as |dp|/g * aire from the hybrid coefficients otherwise. """ if 'masse' in self.ds.variables: return read_array(self.ds, 'masse') pressure = self.interface_pressures() thickness = np.abs(np.diff(pressure, axis=1))/self.gravity return thickness*self.aire[np.newaxis, np.newaxis, :, :] def interface_pressures(self): """ Pressure at the layer interfaces [Pa], shaped (Time, altitude+1, lat, lon), from p = ap + bp*ps. Index 0 is the ground and the last one the top of the atmosphere, where ap = bp = 0. """ ap = read_array(self.ds, 'ap') bp = read_array(self.ds, 'bp') return (ap[np.newaxis, :, np.newaxis, np.newaxis] + bp[np.newaxis, :, np.newaxis, np.newaxis]*self.ps[:, np.newaxis, :, :]) def total(self, field): """ Sum a (Time, altitude, lat, lon) field over the whole planet, one value per time record, leaving out the duplicated longitude column. """ return field[..., :self.nlon].sum(axis=(1, 2, 3)) def mass_of(self, mixing_ratio, record): """ The mass [kg] a mass mixing ratio [kg/kg] amounts to, for one time record: the field weighted by the air mass of every cell. """ return float((mixing_ratio[..., :self.nlon] * self.masse[record][..., :self.nlon]).sum()) def exner(self, exner_type): """ The Exner function pk [J.kg-1.K-1] at the middle of the layers, shaped (Time, altitude, lat, lon). Mars and the generic model run with disvert_type = 2, hence pressure_exner = .false. (iniconst.F90), so "milieu" is the default: it is the form the model itself integrates. "hyb" is the recursive Earth-type form of exner_hyb_m.F90. """ pressure = self.interface_pressures() kappa, cpp, preff = self.kappa, self.cpp, self.preff if None in (kappa, cpp, preff): fail(f"\"{self.path}\" gives no kappa/cpp/preff in " f"\"{CONTROL_NAME}\": the temperature cannot be rebuilt!") nlev = pressure.shape[1] - 1 if nlev < 3: fail("the Exner function needs at least 3 layers!") pk = np.empty_like(pressure[:, :nlev, :, :]) if exner_type == 'milieu': # exner_milieu_m.F90: pk(l) = cpp*((p(l) + p(l+1))/(2*preff))**kappa dum1 = cpp*(2.*preff)**(-kappa) pk[:, :nlev - 1] = dum1*(pressure[:, :nlev - 1] + pressure[:, 1:nlev])**kappa pk[:, nlev - 1] = pk[:, nlev - 2]**2/pk[:, nlev - 3] else: # exner_hyb_m.F90: a downward recursion on alpha and beta, then an # upward one on pk itself, started from the ground value pks pks = cpp*(self.ps/preff)**kappa unpl2k = 1. + 2.*kappa alpha = np.zeros_like(pk) beta = np.zeros_like(pk) beta[:, nlev - 1] = 1./unpl2k for lev in range(nlev - 2, 0, -1): dellta = (pressure[:, lev]*unpl2k + pressure[:, lev + 1]*(beta[:, lev + 1] - unpl2k)) alpha[:, lev] = -pressure[:, lev + 1]/dellta*alpha[:, lev + 1] beta[:, lev] = pressure[:, lev]/dellta pk[:, 0] = ((pressure[:, 0]*pks - 0.5*alpha[:, 1]*pressure[:, 1]) / (pressure[:, 0]*(1. + kappa) + 0.5*(beta[:, 1] - unpl2k)*pressure[:, 1])) for lev in range(1, nlev): pk[:, lev] = alpha[:, lev] + beta[:, lev]*pk[:, lev - 1] return pk def temperature_and_pressure(self, exner_type): """ Temperature [K] and pressure [Pa] at the middle of the layers, rebuilt from "teta" the way calfis.F does it for the physics: T = teta*pk/cpp and pplay = preff*(pk/cpp)**(1/kappa). """ pk = self.exner(exner_type) temperature = read_array(self.ds, 'teta')*pk/self.cpp pressure = self.preff*(pk/self.cpp)**(1./self.kappa) return temperature, pressure def tracers(self): """ The names of the variables carried on the tracer grid, in file order. """ shape = ('Time', 'altitude', 'latitude', 'longitude') return [name for name, var in self.ds.variables.items() if var.dimensions == shape and name not in NOT_TRACERS] def water_saturation(temperature, pressure): """ Water mass mixing ratio at saturation [kg/kg], transcribed unchanged from watersat_mod.F90 so that the cap matches what the model itself enforces. """ epsi = 18./44. psat = 100.*10**(2.07023 - 0.00320991*temperature - 2484.896/temperature + 3.56654*np.log10(temperature)) # Above saturation of the air itself the ratio is meaningless and the model # sets it to 1; the guarded denominator only keeps that branch from warning saturated = psat > pressure denominator = np.where(saturated, 1., pressure - (1. - epsi)*psat) return np.where(saturated, 1., np.maximum(epsi*psat/denominator, 1.e-30)) def molar_mass(names, override): """ The molar mass [kg.mol-1] of a group of tracers, from --molar-mass or from the species their names start with. All the tracers of a group must be the same species, since one number has to stand for the whole group. """ if override is not None: return override species = set() for name in names: head = name.split('_')[0].lower() if head not in MOLAR_MASSES: fail(f"no molar mass known for \"{name}\": give it with --molar-mass!") species.add(head) if len(species) > 1: fail(f"the tracers {', '.join(names)} are not the same species: " f"give one molar mass with --molar-mass!") return MOLAR_MASSES[species.pop()] class Metric: """ The conversions between a tracer mass and the units a target may be given in, for one time record of one file. """ def __init__(self, grid, record, names, override): self.grid = grid self.area = grid.area_total self.air_mass = float(grid.air_mass[record]) self.names = names self.override = override self._molar_mass = None def molar_mass(self): """ The molar mass of the group, resolved on first use only, so that a run that never asks for mol/mol never needs one. """ if self._molar_mass is None: if self.grid.molar_mass_air is None: fail(f"\"{self.grid.path}\" gives no kappa/cpp in " f"\"{CONTROL_NAME}\": mol/mol cannot be converted!") self._molar_mass = molar_mass(self.names, self.override) return self._molar_mass def from_mass(self, mass, unit): """ Convert a tracer mass [kg] into 'unit'. """ if unit == 'kg': return mass if unit == 'kg/m2': return mass/self.area if unit == 'pr-um': return mass/self.area/RHO_WATER*1.e6 if unit == 'kg/kg': return mass/self.air_mass if unit == 'mol/mol': return mass/self.air_mass*self.grid.molar_mass_air/self.molar_mass() fail(f"unknown unit \"{unit}\"!") def to_mass(self, value, unit): """ Convert a value given in 'unit' into a tracer mass [kg]. """ if unit == 'kg': return value if unit == 'kg/m2': return value*self.area if unit == 'pr-um': return value*self.area*RHO_WATER*1.e-6 if unit == 'kg/kg': return value*self.air_mass if unit == 'mol/mol': return value*self.air_mass*self.molar_mass()/self.grid.molar_mass_air fail(f"unknown unit \"{unit}\"!") def report(self, label, mass, unit=None): """ Print a tracer mass in every unit, or in one unit only when 'unit' is given (mol/mol is skipped unless asked for, as it needs a molar mass). """ units = [unit] if unit else [u for u in UNITS if u != 'mol/mol'] print(f" {label}") for name in units: print(f" {self.from_mass(mass, name):>16.8e} {name}") def group_mass(grid, fields, record): """ The total mass [kg] of a group of tracer fields for one time record. """ return sum(grid.mass_of(field[record], record) for field in fields.values()) def capped_mass(grid, fields, factor, record, qsat): """ The mass [kg] the group would hold once scaled by 'factor', with the saturating tracers clipped at qsat. Monotonic in 'factor', which is what lets a bisection solve for the target. """ total = 0. for name, field in fields.items(): scaled = factor*field[record] if qsat is not None and name in SATURATING: scaled = np.minimum(scaled, qsat[record]) total += grid.mass_of(scaled, record) return total def solve_factor(grid, fields, record, target_mass, qsat, metric, unit): """ The factor bringing the group to 'target_mass', by bisection on the mass the group holds once capped at saturation. Without a cap this is just the ratio of the target to the current inventory. """ current = group_mass(grid, fields, record) if qsat is None: if current <= 0.: fail(f"the current inventory is {current:.6e} kg: there is nothing " f"to scale up!") return target_mass/current, 0 # Everything saturating is bounded from above; anything else is not, so the # reachable maximum is only finite when the whole group saturates if all(name in SATURATING for name in fields): reachable = grid.mass_of(qsat[record], record) if target_mass > reachable: fail(f"the target is out of reach: a fully saturated atmosphere " f"holds at most {metric.from_mass(reachable, unit):.8e} {unit}!") # Bracket the solution, starting from the factor an uncapped run would use low, high = 0., max(target_mass/current, 1.) if current > 0. else 1. for _ in range(200): if capped_mass(grid, fields, high, record, qsat) >= target_mass: break high *= 2. else: fail("the target could not be bracketed: check it is reachable!") for _ in range(200): middle = 0.5*(low + high) if capped_mass(grid, fields, middle, record, qsat) < target_mass: low = middle else: high = middle if high - low <= 1.e-14*high: break factor = 0.5*(low + high) capped = 0 for name, field in fields.items(): if name in SATURATING: capped += int(np.count_nonzero( factor*field[record][..., :grid.nlon] > qsat[record][..., :grid.nlon])) return factor, capped class Reservoir: """ The surface reservoir of a startfi.nc file, and the mass it holds. The physics grid is unstructured and split into sub-slopes: the physical area of a sub-slope is area*subslope_dist/cos(slope), exactly the weight slope_weight returns in PEM/src/common/slopes.F90. """ def __init__(self, ds, name, grid): self.ds = ds self.path = ds.filepath() self.name = name if name not in ds.variables: fail(f"\"{self.path}\" has no surface reservoir \"{name}\"!") self.var = ds.variables[name] area = read_array(ds, 'area') self.ngrid = area.size expected = (grid.nlat - 2)*grid.nlon + 2 if self.ngrid != expected: fail(f"\"{self.path}\" has {self.ngrid} physical points but " f"\"{grid.path}\" needs {expected}: the two grids do not match!") self.qsurf = read_array(ds, name) if self.qsurf.ndim != 3: fail(f"\"{name}\" of \"{self.path}\" is not a " f"(Time, nslope, physical_points) field!") self.ntime, self.nslope = self.qsurf.shape[0], self.qsurf.shape[1] self.weight = self.slope_weight(grid)*area[np.newaxis, :] def slope_weight(self, grid): """ The weight of every sub-slope of every grid point, shaped (nslope, ngrid). Falls back to a flat unit distribution when the file carries no sub-slope information at all. """ if 'subslope_dist' not in self.ds.variables: warn(f"\"{self.path}\" has no \"subslope_dist\": the sub-slopes are " f"taken as flat and evenly distributed.") return np.full((self.nslope, self.ngrid), 1./self.nslope) dist = read_array(self.ds, 'subslope_dist') if dist.shape == (self.nslope, self.ngrid): pass elif dist.shape == (self.ngrid, self.nslope): dist = dist.T elif dist.size == self.nslope*self.ngrid: dist = dist.reshape(self.nslope, self.ngrid) else: fail(f"\"subslope_dist\" of \"{self.path}\" has shape {dist.shape}, " f"which is neither ({self.nslope}, {self.ngrid}) nor its " f"transpose!") if 'def_slope' not in self.ds.variables: warn(f"\"{self.path}\" has no \"def_slope\": the sub-slopes are " f"taken as flat.") return dist edges = read_array(self.ds, 'def_slope') if edges.size != self.nslope + 1: fail(f"\"def_slope\" of \"{self.path}\" holds {edges.size} bounds " f"for {self.nslope} sub-slopes!") # Mean slope of each bin, as define_slopes computes it mean = 0.5*(edges[:-1] + edges[1:]) return dist/np.cos(np.radians(mean))[:, np.newaxis] def mass(self, record): """ The mass [kg] the reservoir holds for one time record. """ return float((self.qsurf[record]*self.weight).sum()) def draw(self, record, amount): """ Take 'amount' [kg] out of the reservoir (give it back when negative), proportionally to what each point already holds so that the ice only moves where there is ice. Returns the updated field. """ held = self.mass(record) if held <= 0.: fail(f"the reservoir \"{self.name}\" holds {held:.6e} kg: there is " f"nothing to draw on and nowhere to put anything back!") if amount > held: fail(f"the reservoir \"{self.name}\" holds {held:.6e} kg but " f"{amount:.6e} kg would have to be taken out of it!") return self.qsurf[record]*(held - amount)/held def inventory(path, args): """ Report the inventory of every requested tracer without touching anything. """ with Dataset(path, mode='r') as ds: grid = Grid(ds) names = args.tracer.split(',') if args.tracer else grid.tracers() if not names: fail(f"\"{path}\" holds no tracer!") describe(grid, path) for record in range(grid.ntime): metric = Metric(grid, record, names, args.molar_mass) print(f"\nRecord {record + 1}/{grid.ntime}:") for name in names: if name not in ds.variables: fail(f"\"{path}\" has no tracer \"{name}\"! Available: " f"{', '.join(grid.tracers())}") mass = group_mass(grid, {name: read_array(ds, name)}, record) metric.report(name, mass, args.unit) def describe(grid, path): """ Print what the file says about its own planet and grid. """ print(f"\nFile: \"{path}\"") print(f" Grid {grid.nlon} x {grid.nlat - 1} x {grid.nlev} " f"(lon x lat x alt), {grid.ntime} time record(s)") print(f" Gravity {grid.gravity:>16.8e} m.s-2") if grid.radius: sphere = 4.*np.pi*grid.radius**2 print(f" Planet radius {grid.radius:>16.8e} m") print(f" Total area {grid.area_total:>16.8e} m2 " f"({grid.area_total/sphere:.6f} x 4*pi*R2)") else: print(f" Total area {grid.area_total:>16.8e} m2") if grid.molar_mass_air: print(f" Air molar mass {grid.molar_mass_air:>16.8e} kg.mol-1") for record in range(grid.ntime): print(f" Air mass (record {record + 1}) {grid.air_mass[record]:>16.8e} kg") def solve(path, args): """ Work out, without writing anything, what the rescaling amounts to: the new tracer fields, and the new surface reservoir when a startfi.nc is given. Nothing is copied or opened for writing until this has gone through, so a target that turns out to be out of reach leaves no half-made file behind. """ names = args.tracer.split(',') reservoir = None fi = None with Dataset(path, mode='r') as ds: grid = Grid(ds) describe(grid, path) available = grid.tracers() for name in names: if name not in ds.variables: fail(f"\"{path}\" has no tracer \"{name}\"! Available: " f"{', '.join(available)}") fields = {name: read_array(ds, name) for name in names} qsat = None if args.saturation: if not any(name in SATURATING for name in names): warn(f"--saturation caps only {', '.join(SATURATING)}, none of " f"which is being rescaled: nothing will be clipped.") else: temperature, pressure = grid.temperature_and_pressure(args.exner) qsat = water_saturation(temperature, pressure) try: if args.startfi: fi = Dataset(args.startfi, mode='r') reservoir = Reservoir(fi, args.reservoir, grid) if reservoir.ntime != grid.ntime: fail(f"\"{args.startfi}\" has {reservoir.ntime} time " f"record(s) but \"{path}\" has {grid.ntime}!") for record in range(grid.ntime): metric = Metric(grid, record, names, args.molar_mass) before = group_mass(grid, fields, record) target_mass = metric.to_mass(args.target, args.unit) factor, capped = solve_factor(grid, fields, record, target_mass, qsat, metric, args.unit) print(f"\nRecord {record + 1}/{grid.ntime}: {' + '.join(names)}") metric.report("Before", before, args.unit) for name, field in fields.items(): scaled = factor*field[record] if qsat is not None and name in SATURATING: scaled = np.minimum(scaled, qsat[record]) field[record] = scaled after = group_mass(grid, fields, record) metric.report("After", after, args.unit) print(f" factor applied {factor:>16.8e}") if qsat is not None: print(f" cells capped at saturation: {capped}") if reservoir is not None: held = reservoir.mass(record) reservoir.qsurf[record] = reservoir.draw(record, after - before) print(f" Budget with \"{reservoir.name}\" of " f"\"{args.startfi}\"") print(f" surface before {held:>16.8e} kg") print(f" surface after {reservoir.mass(record):>16.8e} kg") print(f" total before {before + held:>16.8e} kg") print(f" total after {after + reservoir.mass(record):>16.8e} kg") finally: if fi is not None: fi.close() return fields, reservoir def write(path, args, fields, reservoir): """ Copy the input files to their outputs (or warn that they are edited in place) and put the new fields in, keeping the type, shape and attributes each variable already has. """ output = path if args.inplace else args.output if args.inplace: warn(f"\"{path}\" is going to be modified in place!") else: shutil.copy2(path, output) with Dataset(output, mode='r+') as ds: for name, field in fields.items(): ds.variables[name][...] = field print(f"\nWritten: \"{output}\"") if reservoir is None: return startfi_out = args.startfi if args.inplace else args.startfi_output if args.inplace: warn(f"\"{args.startfi}\" is going to be modified in place!") else: shutil.copy2(args.startfi, startfi_out) with Dataset(startfi_out, mode='r+') as fi: fi.variables[reservoir.name][...] = reservoir.qsurf print(f"Written: \"{startfi_out}\"") def parse_args(): parser = argparse.ArgumentParser( description="Rescale a tracer of a dynamics start file to a target " "global inventory, keeping the shape of its distribution." ) parser.add_argument( "start", nargs="?", help="Path to the start file (if omitted, you'll be prompted)" ) parser.add_argument( "--tracer", type=str, help="Tracer to rescale, or several separated by commas to set their " "summed inventory with a common factor, e.g. 'h2o_vap,h2o_ice'. " "Without it, every tracer of the file is reported" ) parser.add_argument( "--target", type=float, help="Target inventory. Without it, the current one is reported and " "nothing is written" ) parser.add_argument( "--unit", type=str, choices=UNITS, help="Unit of the target: total mass [kg], global mean column [kg/m2] " "or its water equivalent [pr-um], global mean mass mixing ratio " "[kg/kg] or volume mixing ratio [mol/mol] (default: pr-um). " "Without a target it selects the single unit reported" ) parser.add_argument( "--output", type=str, help="Output file (default: 'restart.nc' next to the input, the name a " "run gives to the state it produces)" ) parser.add_argument( "--inplace", action="store_true", help="Modify the input file itself. Be careful: your input file is " "modified!" ) parser.add_argument( "--force", action="store_true", help="Overwrite the output file when it already exists. Without it an " "existing 'restart.nc' left by a run is not silently replaced" ) parser.add_argument( "--saturation", action="store_true", help="Keep the water vapour below saturation: the field is clipped at " "qsat and the factor solved so that the target is still met" ) parser.add_argument( "--exner", type=str, default="milieu", choices=("milieu", "hyb"), help="Form of the Exner function used to rebuild the temperature from " "'teta'. 'milieu' is what Mars and the generic model integrate " "(disvert_type = 2); 'hyb' is the Earth-type one (default: milieu)" ) parser.add_argument( "--startfi", type=str, help="Companion startfi.nc whose surface reservoir absorbs the mass " "added to (or removed from) the atmosphere, so that the total " "inventory is conserved" ) parser.add_argument( "--startfi-output", type=str, help="Output file for --startfi (default: 'restartfi.nc' next to it)" ) parser.add_argument( "--reservoir", type=str, help="Surface reservoir to draw on in --startfi (default: the ice " "tracer of the group, e.g. 'h2o_ice')" ) parser.add_argument( "--molar-mass", type=float, dest="molar_mass", help="Molar mass of the tracer [kg.mol-1], for the 'mol/mol' unit. " "Taken from the tracer name when it is a known species" ) return parser.parse_args() def default_output(path): """ Where the result goes when no output name is given, next to its input: the convention of a run, which reads start.nc and startfi.nc and writes restart.nc and restartfi.nc. A file named otherwise falls back to its own name with '_rescaled' appended. """ folder, name = os.path.split(path) if name.startswith('start'): return os.path.join(folder, f"re{name}") stem, extension = os.path.splitext(name) return os.path.join(folder, f"{stem}_rescaled{extension}") def check_output(path, force): """ Refuse to replace an output file that is already there, unless --force says so: the default name is the one a run gives to its own restart file. """ if os.path.exists(path) and not force: fail(f"\"{path}\" already exists: give another --output, or --force to " f"replace it!") def main(): args = parse_args() path = args.start if not path: # Interactive mode: enable tab completion for filenames readline.set_completer(complete_filename) readline.parse_and_bind("tab: complete") try: path = input("Enter the path to the start file: ").strip() except (EOFError, KeyboardInterrupt): print("\nExiting.") return if not path: print("No file specified. Exiting.") return if not os.path.isfile(path): fail(f"file \"{path}\" not found!") if args.target is None: if args.output or args.inplace or args.startfi: warn("no --target given: nothing is written.") inventory(path, args) print("\nDone!") return if not args.tracer: fail("--target needs --tracer to say what is being rescaled!") if not args.unit: args.unit = 'pr-um' if args.output and args.inplace: fail("--output and --inplace cannot be used together!") if not args.output: args.output = default_output(path) if not args.inplace: if os.path.abspath(args.output) == os.path.abspath(path): fail("the output file is the input file: use --inplace to say so!") check_output(args.output, args.force) if args.startfi: if not os.path.isfile(args.startfi): fail(f"file \"{args.startfi}\" not found!") if args.startfi_output and args.inplace: fail("--startfi-output and --inplace cannot be used together!") if not args.startfi_output: args.startfi_output = default_output(args.startfi) if not args.inplace: if os.path.abspath(args.startfi_output) == os.path.abspath(args.startfi): fail("the startfi output file is the startfi input file: " "use --inplace to say so!") check_output(args.startfi_output, args.force) if not args.reservoir: # The condensed form of the species being rescaled is the reservoir # it exchanges with: h2o_vap, or h2o_vap + h2o_ice, draws on h2o_ice species = {name.split('_')[0] for name in args.tracer.split(',')} if len(species) != 1: fail("--startfi needs --reservoir to say which surface " "reservoir absorbs the change!") args.reservoir = f"{species.pop()}_ice" elif args.startfi_output or args.reservoir: warn("--startfi-output and --reservoir do nothing without --startfi.") fields, reservoir = solve(path, args) write(path, args, fields, reservoir) print("Done!") if __name__ == "__main__": main()