From 5e37ccca720b7a8b8b3fe16011dc65966b73c29b Mon Sep 17 00:00:00 2001 From: Dietmar Muller Date: Thu, 11 Jun 2026 22:54:08 +0000 Subject: [PATCH 1/6] Replace PlateTectonicTools dependency with GPlately PlateTectonicTools is now fully integrated into GPlately (as the 'gplately.ptt' module) so optAPM now depends on GPlately instead. Note that GPlately's ContinentContouring.calculate_contoured_continents() expects each continent polygon to be a 2-tuple (polygon, reconstructed feature geometry), unlike the old PTT version (which accepted plain polygons), so the caller in plate_velocity_partitioner.py was updated accordingly. Also updated the 'ptt' imports in the figures notebooks. --- continent_fragmentation.py | 2 +- figures/generate_plate_velocity_figures.ipynb | 4 ++-- figures/generate_trench_migration_figures.ipynb | 4 ++-- optimised_rotation_updater.py | 2 +- plate_velocity_partitioner.py | 7 ++++++- 5 files changed, 12 insertions(+), 7 deletions(-) diff --git a/continent_fragmentation.py b/continent_fragmentation.py index 34d08a3..71e3a66 100644 --- a/continent_fragmentation.py +++ b/continent_fragmentation.py @@ -2,7 +2,7 @@ import math import os import os.path -from ptt.continent_contours import ContouredContinent, ContinentContouring +from gplately.ptt.continent_contours import ContouredContinent, ContinentContouring import pygplates import sys import numpy as np diff --git a/figures/generate_plate_velocity_figures.ipynb b/figures/generate_plate_velocity_figures.ipynb index c80afd3..c0416c7 100644 --- a/figures/generate_plate_velocity_figures.ipynb +++ b/figures/generate_plate_velocity_figures.ipynb @@ -25,7 +25,7 @@ "\n", "from astropy_healpix import healpy\n", "import pygplates\n", - "from ptt.utils import points_in_polygons\n", + "from gplately.ptt.utils import points_in_polygons\n", "import math\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", @@ -759,4 +759,4 @@ }, "nbformat": 4, "nbformat_minor": 4 -} +} \ No newline at end of file diff --git a/figures/generate_trench_migration_figures.ipynb b/figures/generate_trench_migration_figures.ipynb index f9eaada..7a6d8fe 100644 --- a/figures/generate_trench_migration_figures.ipynb +++ b/figures/generate_trench_migration_figures.ipynb @@ -26,7 +26,7 @@ "from __future__ import print_function\n", "\n", "import pygplates\n", - "from ptt import subduction_convergence as sc\n", + "from gplately.ptt import subduction_convergence as sc\n", "import numpy as np\n", "import os.path\n", "import pandas as pd\n", @@ -576,4 +576,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/optimised_rotation_updater.py b/optimised_rotation_updater.py index bdbe8cc..9ea9f6d 100644 --- a/optimised_rotation_updater.py +++ b/optimised_rotation_updater.py @@ -4,7 +4,7 @@ import sys import warnings -from ptt import remove_plate_rotations +from gplately.ptt import remove_plate_rotations class OptimisedRotationUpdater(object): diff --git a/plate_velocity_partitioner.py b/plate_velocity_partitioner.py index 36fff62..d47e4c0 100644 --- a/plate_velocity_partitioner.py +++ b/plate_velocity_partitioner.py @@ -202,7 +202,12 @@ def generate_points_and_plate_ids( # # NOTE: We are NOT excluding contours based on perimeter/area ratio. # That determination must be made by the final cost function that calculates the cost (penalty). - contoured_continents = self.continent_fragmentation.calculate_contoured_continents(reconstructed_continent_polygons, ref_rotation_start_age) + # + # Note: GPlately's continent contouring (which replaced PlateTectonicTools') expects each + # continent polygon to be a 2-tuple (polygon, reconstructed feature geometry). + contoured_continents = self.continent_fragmentation.calculate_contoured_continents( + list(zip(reconstructed_continent_polygons, reconstructed_feature_geometries)), + ref_rotation_start_age) # Find the contoured continent (if any) containing each continent point. continent_point_contoured_continent_indices = [None] * len(continent_points) From 3e6dede3d9997719921b3bd54a644d0a8f40d9c7 Mon Sep 17 00:00:00 2001 From: Dietmar Muller Date: Thu, 11 Jun 2026 22:54:28 +0000 Subject: [PATCH 2/6] Speed up the optimisation workflow (~10x) via seed screening and other optimisations Changes (an empirical seed-count convergence study found the per-timestep multistart to be heavily over-sampled - see 'seed_study.py' and the README): * Seed screening ('screen then polish'): the objective function is first evaluated once at every seed (costs roughly the same as fully optimising one or two seeds), then only the 'seed_screen_top_n' best seeds plus 'seed_screen_uniform_n' uniformly spread seeds are fully optimised with NLopt. Defaults (16+16 out of 400 seeds) reproduce the full 400-seed multistart minimum to within ~0.5% cost at ~10x less CPU. Set seed_screen_top_n=None to restore the original behaviour. * Cache the ObjectiveFunction per process per timestep (it was previously constructed - including loading rotation/trench/velocity files - once per seed). * Skip the isochron interpolation ('isopolate') and the loading of the ridge/isochron/isocob files when fracture zones are disabled (they are disabled in all current configs; this work was previously done at every timestep and then unused). * Safety cap ('nlopt_max_eval_safety', default 1000) on COBYLA objective evaluations: COBYLA occasionally fails to converge and can cycle for thousands of evaluations, stalling an entire MPI rank and hence the whole timestep. * New resumable study/profiling harness 'seed_study.py' (prep/probe/run/ report) to re-run the seed-count convergence study on any model. * Environment variable overrides (OPTAPM_START_AGE, OPTAPM_END_AGE, OPTAPM_MODELS, OPTAPM_SCREEN_TOP_N, OPTAPM_SCREEN_UNIFORM_N, OPTAPM_SERIAL) for quick test runs without editing the config. --- Optimised_APM.py | 186 ++++++++++++-- Optimised_config.py | 48 ++++ README.md | 59 +++-- optapm.py | 140 ++++++----- seed_study.py | 588 ++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 919 insertions(+), 102 deletions(-) create mode 100644 seed_study.py diff --git a/Optimised_APM.py b/Optimised_APM.py index ccee971..c0bcab8 100644 --- a/Optimised_APM.py +++ b/Optimised_APM.py @@ -347,8 +347,13 @@ def excepthook(t,v,tb): # -------------------------------------------------------------------- # Load all data - data = ms.dataLoader(datadir, rotfile, ref_rotation_file, tm_file=tm_file, pv_file=pv_file, nnr_rotfile=nnr_rotfile, - ridge_file=ridge_file, isochron_file=isochron_file, isocob_file=isocob_file, + # + # OPTIMISATION: The ridge/isochron/isocob files are only used by the fracture zone + # component, so don't load them when fracture zones are disabled (saves time at every timestep). + data = ms.dataLoader(datadir, rotfile, ref_rotation_file, tm_file=tm_file, pv_file=pv_file, nnr_rotfile=nnr_rotfile, + ridge_file=ridge_file if enable_fracture_zones else None, + isochron_file=isochron_file if enable_fracture_zones else None, + isocob_file=isocob_file if enable_fracture_zones else None, hst_file=hst_file, hs_file=hs_file, interpolated_hotspots=interpolated_hotspots) # Calculate starting conditions @@ -451,12 +456,13 @@ def excepthook(t,v,tb): # -------------------------------------------------------------------- # Function to run optimisation routine - def run_optimisation(x, opt_n, N, lb, ub, model_stop_condition, max_iter, interval, rotation_file, - no_net_rotation_file, ref_rotation_start_age, Lats, Lons, spreading_directions, + def run_optimisation(x, opt_n, N, lb, ub, model_stop_condition, max_iter, interval, rotation_file, + no_net_rotation_file, ref_rotation_start_age, Lats, Lons, spreading_directions, spreading_asymmetries, seafloor_ages, PID, CPID, data_array, weights_array, cost_func_code_string_array, bounds_array, trench_migration_file, plate_velocity_file, ref_rotation_end_age, ref_rotation_plate_id, - reformArray, trail_data, use_trail_age_uncertainty, trail_age_uncertainty_ellipse, tm_method): + reformArray, trail_data, use_trail_age_uncertainty, trail_age_uncertainty_ellipse, tm_method, + max_eval_safety=None): # Make sure remote nodes/cores also import these modules (when running code in parallel). # @@ -473,19 +479,29 @@ def run_optimisation(x, opt_n, N, lb, ub, model_stop_condition, max_iter, interv import sys import types - # Turn cost function code strings back into functions. - cost_func_array = [types.FunctionType(marshal.loads(cost_func_code_string), globals(), 'cost_func') - for cost_func_code_string in cost_func_code_string_array] - - # Load up the object function object once (eg, load rotation files). - # NLopt will then call it multiple times. - # NLopt will call this as 'obj_f(x, grad)' because 'obj_f' has a '__call__' method. - obj_f = ObjectiveFunction( - interval, rotation_file, no_net_rotation_file, ref_rotation_start_age, Lats, Lons, spreading_directions, - spreading_asymmetries, seafloor_ages, PID, CPID, data_array, weights_array, cost_func_array, bounds_array, - trench_migration_file, plate_velocity_file, ref_rotation_end_age, ref_rotation_plate_id, reformArray, trail_data, - use_trail_age_uncertainty, trail_age_uncertainty_ellipse, tm_method) - + # OPTIMISATION: Cache the ObjectiveFunction in the current process so that + # constructing it (which loads rotation/trench/velocity files) only happens once + # per process per timestep, instead of once per seed. + obj_f_cache = globals().setdefault('_optapm_objective_function_cache', {}) + obj_f_cache_key = (rotation_file, ref_rotation_start_age, ref_rotation_end_age) + obj_f = obj_f_cache.get(obj_f_cache_key) + if obj_f is None: + # Turn cost function code strings back into functions. + cost_func_array = [types.FunctionType(marshal.loads(cost_func_code_string), globals(), 'cost_func') + for cost_func_code_string in cost_func_code_string_array] + + # Load up the object function object once (eg, load rotation files). + # NLopt will then call it multiple times. + # NLopt will call this as 'obj_f(x, grad)' because 'obj_f' has a '__call__' method. + obj_f = ObjectiveFunction( + interval, rotation_file, no_net_rotation_file, ref_rotation_start_age, Lats, Lons, spreading_directions, + spreading_asymmetries, seafloor_ages, PID, CPID, data_array, weights_array, cost_func_array, bounds_array, + trench_migration_file, plate_velocity_file, ref_rotation_end_age, ref_rotation_plate_id, reformArray, trail_data, + use_trail_age_uncertainty, trail_age_uncertainty_ellipse, tm_method) + # Only keep the cached objective function for the current timestep. + obj_f_cache.clear() + obj_f_cache[obj_f_cache_key] = obj_f + opt = nlopt.opt(nlopt.LN_COBYLA, opt_n) opt.set_min_objective(obj_f) opt.set_lower_bounds(lb) @@ -501,6 +517,13 @@ def run_optimisation(x, opt_n, N, lb, ub, model_stop_condition, max_iter, interv opt.set_ftol_rel(1e-6) opt.set_xtol_rel(1e-8) + # Safety cap: COBYLA occasionally fails to converge (it can cycle for + # thousands of evaluations, stalling an entire process and hence the whole + # timestep). The vast majority of optimisations converge in well under + # 200 evaluations, so a generous cap only affects pathological cases. + if max_eval_safety: + opt.set_maxeval(max_eval_safety) + xopt = opt.optimize(x) minf = opt.last_optimum_value() @@ -526,16 +549,137 @@ def run_optimisation(x, opt_n, N, lb, ub, model_stop_condition, max_iter, interv # Start optimisation # Wrap 'run_optimisation()' by passing all the constant parameters (ie, everything except 'x'). - runopt = partial(run_optimisation, opt_n=opt_n, N=N, lb=lb, ub=ub, + runopt = partial(run_optimisation, opt_n=opt_n, N=N, lb=lb, ub=ub, model_stop_condition=model_stop_condition, max_iter=max_iter, interval=interval, rotation_file=rotation_file, - no_net_rotation_file=no_net_rotation_file, ref_rotation_start_age=ref_rotation_start_age, + no_net_rotation_file=no_net_rotation_file, ref_rotation_start_age=ref_rotation_start_age, Lats=Lats, Lons=Lons, spreading_directions=spreading_directions, spreading_asymmetries=spreading_asymmetries, seafloor_ages=seafloor_ages, PID=PID, CPID=CPID, data_array=data_array, weights_array=weights_array, cost_func_code_string_array=cost_func_code_string_array, bounds_array=bounds_array, trench_migration_file=trench_migration_file, plate_velocity_file=plate_velocity_file, ref_rotation_end_age=ref_rotation_end_age, ref_rotation_plate_id=ref_rotation_plate_id, reformArray=reformArray, trail_data=trail_data, use_trail_age_uncertainty=use_trail_age_uncertainty, - trail_age_uncertainty_ellipse=trail_age_uncertainty_ellipse, tm_method=tm_method) + trail_age_uncertainty_ellipse=trail_age_uncertainty_ellipse, tm_method=tm_method, + max_eval_safety=nlopt_max_eval_safety) + + + # -------------------------------------------------------------------- + # -------------------------------------------------------------------- + # Optionally screen the seeds (a single objective function evaluation per seed) + # and then only fully optimise the most promising ones ("screen then polish"). + # + # A full NLopt optimisation typically uses ~80-200 objective evaluations, so + # screening all seeds costs roughly the same as fully optimising one or two seeds. + + if seed_screen_top_n is not None: + + # Function to evaluate the objective function once at a seed. + # Self-contained (all imports inside) so it can run on remote nodes/cores. + def run_seed_screening(x, opt_n, N, lb, ub, model_stop_condition, max_iter, interval, rotation_file, + no_net_rotation_file, ref_rotation_start_age, Lats, Lons, spreading_directions, + spreading_asymmetries, seafloor_ages, PID, CPID, + data_array, weights_array, cost_func_code_string_array, bounds_array, + trench_migration_file, plate_velocity_file, ref_rotation_end_age, ref_rotation_plate_id, + reformArray, trail_data, use_trail_age_uncertainty, trail_age_uncertainty_ellipse, tm_method, + max_eval_safety=None): + + from objective_function import ObjectiveFunction + import marshal + import types + + # Share the per-process ObjectiveFunction cache with 'run_optimisation()'. + obj_f_cache = globals().setdefault('_optapm_objective_function_cache', {}) + obj_f_cache_key = (rotation_file, ref_rotation_start_age, ref_rotation_end_age) + obj_f = obj_f_cache.get(obj_f_cache_key) + if obj_f is None: + cost_func_array = [types.FunctionType(marshal.loads(cost_func_code_string), globals(), 'cost_func') + for cost_func_code_string in cost_func_code_string_array] + obj_f = ObjectiveFunction( + interval, rotation_file, no_net_rotation_file, ref_rotation_start_age, Lats, Lons, spreading_directions, + spreading_asymmetries, seafloor_ages, PID, CPID, data_array, weights_array, cost_func_array, bounds_array, + trench_migration_file, plate_velocity_file, ref_rotation_end_age, ref_rotation_plate_id, reformArray, trail_data, + use_trail_age_uncertainty, trail_age_uncertainty_ellipse, tm_method) + obj_f_cache.clear() + obj_f_cache[obj_f_cache_key] = obj_f + + return obj_f(x, None) + + runscreen = partial(run_seed_screening, opt_n=opt_n, N=N, lb=lb, ub=ub, + model_stop_condition=model_stop_condition, max_iter=max_iter, interval=interval, rotation_file=rotation_file, + no_net_rotation_file=no_net_rotation_file, ref_rotation_start_age=ref_rotation_start_age, + Lats=Lats, Lons=Lons, spreading_directions=spreading_directions, spreading_asymmetries=spreading_asymmetries, + seafloor_ages=seafloor_ages, PID=PID, CPID=CPID, + data_array=data_array, weights_array=weights_array, cost_func_code_string_array=cost_func_code_string_array, bounds_array=bounds_array, + trench_migration_file=trench_migration_file, plate_velocity_file=plate_velocity_file, + ref_rotation_end_age=ref_rotation_end_age, ref_rotation_plate_id=ref_rotation_plate_id, + reformArray=reformArray, trail_data=trail_data, use_trail_age_uncertainty=use_trail_age_uncertainty, + trail_age_uncertainty_ellipse=trail_age_uncertainty_ellipse, tm_method=tm_method) + + # Select the seeds to fully optimise: the 'top_n' best-screened seeds plus + # 'uniform_n' seeds spread uniformly across the search space (the latter is + # insurance in case a poorly-screened seed would have converged to the global + # minimum). Note that the uniform subset is selected by striding the seed list, + # which is in grid order and hence spatially spread out. + def select_polish_seeds(seeds, costs, top_n, uniform_n): + selected_indices = set(np.argsort(costs)[:top_n].tolist()) + if uniform_n: + stride = max(1, len(seeds) // uniform_n) + selected_indices.update(range(0, len(seeds), stride)) + return [seeds[index] for index in sorted(selected_indices)] + + if use_parallel == IPYPARALLEL: + + screen_costs = dview.map(runscreen, x) + x = select_polish_seeds(x, screen_costs, seed_screen_top_n, seed_screen_uniform_n) + print("Screened {0} seeds; fully optimising {1} of them.".format(len(screen_costs), len(x))) + sys.stdout.flush() + + elif use_parallel == MPI4PY: + + # Each process screens the seeds it received from the root process. + screen_costs = [runscreen(x_item) for x_item in x] + + # Gather all screened costs and seeds into the root (0) process. + all_screen_costs = mpi_comm.gather(screen_costs, root=0) + all_screen_seeds = mpi_comm.gather(x, root=0) + + if mpi_rank == 0: + # Flatten the list of lists into a single list. + all_screen_costs = list(itertools.chain.from_iterable(all_screen_costs)) + all_screen_seeds = list(itertools.chain.from_iterable(all_screen_seeds)) + + polish_seeds = select_polish_seeds( + all_screen_seeds, all_screen_costs, seed_screen_top_n, seed_screen_uniform_n) + + print("Screened {0} seeds; fully optimising {1} of them (best {2} plus {3} uniformly spread).".format( + len(all_screen_seeds), len(polish_seeds), seed_screen_top_n, seed_screen_uniform_n)) + sys.stdout.flush() + + # Divide the polish seeds among the processes (same way as the original seeds). + if len(polish_seeds) < mpi_size: + polish_seed_lists = [[polish_seed] for polish_seed in polish_seeds] + polish_seed_lists.extend([[]] * (mpi_size - len(polish_seeds))) + else: + num_seeds_per_rank = len(polish_seeds) // mpi_size + polish_seed_lists = [] + for mpi_index in range(mpi_size): + seed_index = mpi_index * num_seeds_per_rank + polish_seed_lists.append(polish_seeds[seed_index : seed_index + num_seeds_per_rank]) + for seed_index in range(mpi_size * num_seeds_per_rank, len(polish_seeds)): + polish_seed_lists[seed_index - mpi_size * num_seeds_per_rank].append(polish_seeds[seed_index]) + else: + polish_seed_lists = None + + # Scatter the polish seeds across all processes (from root process). + x = mpi_comm.scatter(polish_seed_lists, root=0) + + else: + + # Calculate serially. + screen_costs = [runscreen(x_item) for x_item in x] + x = select_polish_seeds(x, screen_costs, seed_screen_top_n, seed_screen_uniform_n) + print("Screened {0} seeds; fully optimising {1} of them.".format(len(screen_costs), len(x))) + sys.stdout.flush() + # Start timer for current time step. #start = time.time() diff --git a/Optimised_config.py b/Optimised_config.py index 107141c..6a3906b 100644 --- a/Optimised_config.py +++ b/Optimised_config.py @@ -506,6 +506,31 @@ def get_reference_params(age): return ref_rotation_plate_id, ref_rotation_file +# +# Seed screening ("screen then polish"). +# +# The objective function is first evaluated once at every seed (cheap screening), and then +# only the most promising seeds are fully optimised with NLopt. A full NLopt optimisation +# typically uses ~80-200 objective evaluations, so screening all seeds costs roughly the +# same as fully optimising just one or two seeds - while identifying where the low-cost +# basins are. An empirical convergence study (see 'seed_study.py') found that fully +# optimising the best ~16 screened seeds plus a small uniformly-distributed backstop +# reproduces the full 400-seed multistart minimum to within ~0.5% cost at ~10-25x less CPU. +# +# Set 'seed_screen_top_n' to None to disable screening (original behaviour: fully optimise every seed). +seed_screen_top_n = 16 # Fully optimise the N best-screened seeds. +seed_screen_uniform_n = 16 # Also fully optimise this many seeds spread uniformly across the + # search space (insurance in case screening misses a basin - + # eg, when a poor seed would have converged to the global minimum). + +# Safety cap on NLopt objective evaluations per seed (only used when model_stop_condition == 'threshold'). +# COBYLA occasionally fails to converge (it can cycle for thousands of evaluations, stalling +# an entire MPI rank and hence the whole timestep). The vast majority of optimisations +# converge in well under 200 evaluations, so this generous cap only affects pathological cases. +# Set to None to disable the cap. +nlopt_max_eval_safety = 1000 + + search = "Initial" # If True then temporarily expand search radius to 180 whenever the reference plate changes. # Normally the reference plate stays constant at Africa (701), but does switch to 101 for the 1Ga model. @@ -571,6 +596,29 @@ def get_reference_params(age): plot = False +# +# Environment variable overrides (useful for quick test runs without editing this file). +# +# OPTAPM_START_AGE / OPTAPM_END_AGE : Override the age range (eg, to run a single timestep). +# OPTAPM_MODELS : Override the number of seed models. +# OPTAPM_SERIAL=1 : Disable parallelisation (eg, to run/debug on a desktop without MPI). +# +if 'OPTAPM_START_AGE' in os.environ: + start_age = int(os.environ['OPTAPM_START_AGE']) +if 'OPTAPM_END_AGE' in os.environ: + end_age = actual_end_age = int(os.environ['OPTAPM_END_AGE']) +if 'OPTAPM_MODELS' in os.environ: + models = int(os.environ['OPTAPM_MODELS']) +if 'OPTAPM_SCREEN_TOP_N' in os.environ: + seed_screen_top_n = int(os.environ['OPTAPM_SCREEN_TOP_N']) + if seed_screen_top_n < 0: # negative disables screening + seed_screen_top_n = None +if 'OPTAPM_SCREEN_UNIFORM_N' in os.environ: + seed_screen_uniform_n = int(os.environ['OPTAPM_SCREEN_UNIFORM_N']) +if os.environ.get('OPTAPM_SERIAL') == '1': + use_parallel = None + + # # How to handle warnings. # diff --git a/README.md b/README.md index 42cf7fb..fe95837 100644 --- a/README.md +++ b/README.md @@ -15,11 +15,9 @@ The data for the 1Ga model can be obtained [here](https://www.earthbyte.org/webd The following Python packages are required: -* PlateTectonicTools>=0.5 +* gplately>=1.3 - > __Note:__ Until version 0.5 of PlateTectonicTools is available (on conda and pip) you'll need to install the latest version from Github: - > - > `pip install git+https://github.com/EarthByte/PlateTectonicTools` + > __Note:__ PlateTectonicTools is no longer required - its functionality is now fully integrated into [GPlately](https://github.com/GPlates/gplately) (as the `gplately.ptt` module). * pygplates>=1.0 * numpy @@ -44,7 +42,7 @@ First install dependencies that are available on `conda` (in the _conda-forge_ c ``` conda create -n -c conda-forge \ - "pygplates>=1.0" numpy scipy platetectonictools scikit-image pandas xlrd==1.2.0 NLopt \ + "pygplates>=1.0" "gplately>=1.3" numpy scipy scikit-image pandas xlrd==1.2.0 NLopt \ future cartopy matplotlib ipyparallel openpyxl ``` @@ -69,13 +67,6 @@ conda install pip pip install pmagpy ``` -> Note: Until version 0.5 of PlateTectonicTools is available (on conda) you'll need to upgrade it from Github: -> -> ``` -> conda install git -> pip install --upgrade git+https://github.com/EarthByte/PlateTectonicTools -> ``` - ### Install on a HPC system Installation on a High Performance Computing (HPC) system can also be done with a local installation of `conda` (and `pip`). However the exception, compared with installing on desktop, is `mpi4py`. It will likely need to be installed differently to ensure that the MPI implementation of the HPC system is used (instead of conda's MPI). @@ -90,7 +81,7 @@ Similarly to installing on desktop, start by creating a conda environment: ``` conda create -n -c conda-forge \ - "pygplates>=1.0" numpy scipy platetectonictools scikit-image pandas xlrd==1.2.0 NLopt \ + "pygplates>=1.0" "gplately>=1.3" numpy scipy scikit-image pandas xlrd==1.2.0 NLopt \ future cartopy matplotlib ipyparallel openpyxl ``` @@ -121,13 +112,6 @@ Then, similarly to installing on desktop, use `pip` to install `pmagpy` (into th pip install pmagpy ``` -> Note: Until version 0.5 of PlateTectonicTools is available (on conda) you'll need to upgrade it from Github: -> -> ``` -> conda install git -> pip install --upgrade git+https://github.com/EarthByte/PlateTectonicTools -> ``` - #### Job submission script In our example [job submission script](Optimised_APM.sh) (that runs on [NCI's Gadi](https://nci.org.au/our-systems/hpc-systems)) we have the following commands (after the various PBS directives) that: @@ -169,6 +153,35 @@ All settings are now in "Optimised_config.py", such as the model name, start/end Essentially the workflow optimises the absolute plate motion by perturbing the absolute rotation (pole and angle) and then calculating penalties derived from that rotation (eg, net rotation (NR), trench migration (TM) and plate velocity (PV); and a hotspot (HS) penalty for 0-80Ma). Then these penalties are scaled by their weights (eg, NR=1, TM=0.5, PV=0.5) and also internal scaling to make sure each penalty is roughly the same when the weights are all 1.0 (but that’s an implementation detail). Then the penalties are added to give one number. The optimization workflow then perturbs the rotation further to try to reduce that number. It does this many times until it settles on a local minimum (called a seed), and does this for a bunch of seeds to find global minimum. +### Seed screening ("screen then polish") + +By default the workflow now *screens* all seeds before optimising. The objective function is evaluated once at every seed (cheap - roughly the cost of fully optimising just one or two seeds, since a full NLopt optimisation typically uses ~80-200 objective evaluations), and then only the most promising seeds are fully optimised: the `seed_screen_top_n` best-screened seeds plus `seed_screen_uniform_n` seeds spread uniformly across the search space (insurance in case screening misses a basin). Both parameters are in "Optimised_config.py". + +An empirical convergence study on a 1Ga global plate model found that screening 400 seeds and fully optimising the best 16 (plus 16 uniform) reproduces the full 400-seed multistart minimum to within ~0.5% cost, while reducing the per-timestep optimisation cost by roughly a factor of 10. The study also found that the optimised pole is only loosely constrained near the minimum (several near-degenerate minima can differ by 10-20 degrees in pole position while differing by less than 0.5% in cost), so the residual error introduced by screening is well below the intrinsic uncertainty of the objective function itself. + +To reproduce or extend the study (eg, for a different plate model or different seed counts), use the standalone, resumable harness "seed_study.py": + +``` + python seed_study.py prep --age 100 # prepare data for one timestep (re-run until 'PREP DONE') + python seed_study.py probe --age 100 # time a single objective evaluation / full optimisation + python seed_study.py run --age 100 # run the study (resumable; re-run until 'ALL TASKS DONE') + python seed_study.py report --age 100 # summarise results (writes a CSV) +``` + +To disable screening (and fully optimise every seed, as in the original workflow), set `seed_screen_top_n = None`. + +The workflow also caps each NLopt (COBYLA) optimisation at `nlopt_max_eval_safety` objective evaluations (default 1000). COBYLA occasionally fails to converge and can otherwise cycle for thousands of evaluations, stalling an entire MPI rank (and hence the whole timestep, since all ranks must finish before the next timestep begins). + +### Quick test runs + +The age range, seed count, screening parameters and parallelisation can be temporarily overridden via environment variables without editing "Optimised_config.py" - useful for quick smoke tests: + +``` + OPTAPM_START_AGE=5 OPTAPM_END_AGE=0 OPTAPM_MODELS=49 OPTAPM_SERIAL=1 python Optimised_APM.py +``` + +(`OPTAPM_SCREEN_TOP_N` and `OPTAPM_SCREEN_UNIFORM_N` override the screening parameters; a negative `OPTAPM_SCREEN_TOP_N` disables screening.) + ### Plate velocity penalty For the 1Ga model, the plate velocity (PV) penalty currently involves multiplying the plate velocity weight by the median velocity across all continents on the globe: @@ -263,15 +276,15 @@ All these rotation files also comprise the entire optimised rotation model. Note that only those rotations that referenced 000 as their fixed plate will be modified (to include the optimised absolute plate motion). You can also optionally remove plates in the plate hierarchy to make it simpler (eg, plates below 701 are typically removed so that 701 directly references 000). -This can be done using the 'remove_plate_rotations' module in Plate Tectonic Tools ( https://github.com/EarthByte/PlateTectonicTools ) For example: +This can be done using the 'remove_plate_rotations' module in GPlately ( https://github.com/GPlates/gplately ) For example: ``` - python -m ptt.remove_plate_rotations -u -a 0.1 5 -p 70 4 3 2 1 -o removed_70_4_3_2_1_ -- ... + python -m gplately.ptt.remove_plate_rotations -u -a 0.1 5 -p 70 4 3 2 1 -o removed_70_4_3_2_1_ -- ... ``` ...where you replace `...` with all the rotation files in the optimised rotation model. ## Plotting results -After running the workflow and post-processing (although you don't need to run `ptt.remove_plate_rotations` for this), you can plot the +After running the workflow and post-processing (although you don't need to run `gplately.ptt.remove_plate_rotations` for this), you can plot the trench migration stats and net rotation curves for the non-optimised and any optimised models using the Jupyter notebooks in the `figures/` directory. diff --git a/optapm.py b/optapm.py index 76ecfd3..c6e66ca 100644 --- a/optapm.py +++ b/optapm.py @@ -102,20 +102,38 @@ def dataLoader(datadir, rot_file, ref_rot_file=None, tm_file=None, pv_file=None, no_net_rotation_file = datadir + nnr_rotfile + # Note: The ridge/isochron/isocob features are only used by the fracture zone + # component of the objective function. Callers can pass None for these + # files (eg, when fracture zones are disabled) to avoid loading them. if ridge_file: features_ri = pgp.FeatureCollection(datadir + ridge_file) - RidgeFile_subset = pgp.FeatureCollection() + + else: + + features_ri = pgp.FeatureCollection() + + RidgeFile_subset = pgp.FeatureCollection() if isochron_file: features_iso = pgp.FeatureCollection(datadir + isochron_file) - IsochronFile_subset = pgp.FeatureCollection() + + else: + + features_iso = pgp.FeatureCollection() + + IsochronFile_subset = pgp.FeatureCollection() if isocob_file: features_isocob = pgp.FeatureCollection(datadir + isocob_file) - IsoCOBFile_subset = pgp.FeatureCollection() + + else: + + features_isocob = pgp.FeatureCollection() + + IsoCOBFile_subset = pgp.FeatureCollection() if hst_file: @@ -846,40 +864,6 @@ def modelStartConditions(params, data, plot=True): - # Build feature subset lists - for feature in features_iso: - if feature.get_valid_time()[1] <= ref_rotation_start_age: - IsochronFile_subset.add(feature) - - for feature in features_ri: - if feature.get_valid_time()[1] <= ref_rotation_start_age: - RidgeFile_subset.add(feature) - - for feature in features_isocob: - if feature.get_valid_time()[1] <= ref_rotation_start_age: - IsoCOBFile_subset.add(feature) - - - - # This is where intertec is actually called - note that these one line interpolates the isochrons, - # but does not reconstruct them - recon_interpolated_isochrons = [] - - output_features = isopolate.interpolate_isochrons(rotation_model, - [RidgeFile_subset, IsochronFile_subset, IsoCOBFile_subset], - interpolate=range(ref_rotation_end_age, ref_rotation_start_age + 1, interpolation_resolution), - #interpolate=0.01, - tessellate_threshold_radians=math.radians(0.5), - output_scalar_spreading_direction=True, - output_scalar_spreading_rate=True, - output_scalar_spreading_asymmetry=True, - output_scalar_age=True, - print_debug_output=0) - - # Here we do the last step that the old intertec did, reconstructing the interpolated - # isochrons to the selected time. - pgp.reconstruct(output_features, rotation_file, recon_interpolated_isochrons, ref_rotation_start_age, 0) - ## Step2 # Take the coverage that intertec produced, and use it to derive arrays of data more easily analysed using numpy commands @@ -892,40 +876,80 @@ def modelStartConditions(params, data, plot=True): PID = [] CPID = [] - for recon_interpolated_isochron in recon_interpolated_isochrons: + # OPTIMISATION: The interpolated isochrons (via 'isopolate') are only used by the + # fracture zone component of the objective function. Skip this relatively expensive + # step entirely when fracture zones are disabled (data_array[0] is the + # 'enable_fracture_zones' flag), which saves time at every timestep. + if data_array[0]: + + # Build feature subset lists + for feature in features_iso: + if feature.get_valid_time()[1] <= ref_rotation_start_age: + IsochronFile_subset.add(feature) + + for feature in features_ri: + if feature.get_valid_time()[1] <= ref_rotation_start_age: + RidgeFile_subset.add(feature) + + for feature in features_isocob: + if feature.get_valid_time()[1] <= ref_rotation_start_age: + IsoCOBFile_subset.add(feature) + + + + # This is where intertec is actually called - note that these one line interpolates the isochrons, + # but does not reconstruct them + recon_interpolated_isochrons = [] + + output_features = isopolate.interpolate_isochrons(rotation_model, + [RidgeFile_subset, IsochronFile_subset, IsoCOBFile_subset], + interpolate=range(ref_rotation_end_age, ref_rotation_start_age + 1, interpolation_resolution), + #interpolate=0.01, + tessellate_threshold_radians=math.radians(0.5), + output_scalar_spreading_direction=True, + output_scalar_spreading_rate=True, + output_scalar_spreading_asymmetry=True, + output_scalar_age=True, + print_debug_output=0) + + # Here we do the last step that the old intertec did, reconstructing the interpolated + # isochrons to the selected time. + pgp.reconstruct(output_features, rotation_file, recon_interpolated_isochrons, ref_rotation_start_age, 0) + + for recon_interpolated_isochron in recon_interpolated_isochrons: - tmp = recon_interpolated_isochron.get_feature() - tmp2 = tmp.get_geometry(coverage_return=pgp.CoverageReturn.geometry_and_scalars) + tmp = recon_interpolated_isochron.get_feature() + tmp2 = tmp.get_geometry(coverage_return=pgp.CoverageReturn.geometry_and_scalars) - if tmp2: + if tmp2: - coverage_geometry, coverage_scalars = tmp2 - coverage_points = coverage_geometry.get_points() - #tmp3 = coverage_points.get_points() + coverage_geometry, coverage_scalars = tmp2 + coverage_points = coverage_geometry.get_points() + #tmp3 = coverage_points.get_points() - for scalar in coverage_scalars.get(pgp.ScalarType.create_gpml('SpreadingDirection')): + for scalar in coverage_scalars.get(pgp.ScalarType.create_gpml('SpreadingDirection')): - spreading_directions.append(scalar) - #spreading_directions.append(coverage_scalars.get(pgp.ScalarType.create_gpml('SpreadingDirection'))) + spreading_directions.append(scalar) + #spreading_directions.append(coverage_scalars.get(pgp.ScalarType.create_gpml('SpreadingDirection'))) - for scalar in coverage_scalars.get(pgp.ScalarType.create_gpml('SpreadingRate')): + for scalar in coverage_scalars.get(pgp.ScalarType.create_gpml('SpreadingRate')): - spreading_rates.append(scalar) + spreading_rates.append(scalar) - for scalar in coverage_scalars.get(pgp.ScalarType.create_gpml('SpreadingAsymmetry')): + for scalar in coverage_scalars.get(pgp.ScalarType.create_gpml('SpreadingAsymmetry')): - spreading_asymmetries.append(scalar) + spreading_asymmetries.append(scalar) - for scalar in coverage_scalars.get(pgp.ScalarType.create_gpml('Age')): + for scalar in coverage_scalars.get(pgp.ScalarType.create_gpml('Age')): - seafloor_ages.append(scalar) + seafloor_ages.append(scalar) - for point in coverage_points.get_points().to_lat_lon_array(): + for point in coverage_points.get_points().to_lat_lon_array(): - Lats.append(point[0]) - Lons.append(point[1]) - PID.append(tmp.get_reconstruction_plate_id()) - CPID.append(tmp.get_conjugate_plate_id()) + Lats.append(point[0]) + Lons.append(point[1]) + PID.append(tmp.get_reconstruction_plate_id()) + CPID.append(tmp.get_conjugate_plate_id()) # TRENCH MIGRATION # Load pre-computed migration data if needed (net rotation or trench migration) diff --git a/seed_study.py b/seed_study.py new file mode 100644 index 0000000..fc21e3e --- /dev/null +++ b/seed_study.py @@ -0,0 +1,588 @@ +""" +Seed-count convergence study and profiling harness for the optAPM workflow. + +The study answers: how many start seeds (multi-start optimisations) per timestep +are actually needed to reliably find the global minimum of the objective function, +and how much cheaper is a "screen-then-polish" strategy (evaluate the objective +once at every seed, then fully optimise only the best few)? + +The harness is checkpointed/resumable so it can also run in environments with +short execution windows: every command can be re-run repeatedly and will continue +where it left off. + +Usage (from the optAPM directory; config comes from Optimised_config.py): + + # 1. Prepare data for a timestep (re-run until it prints 'PREP DONE'): + python seed_study.py prep --age 100 + + # 2. Time a single objective evaluation and a single full optimisation: + python seed_study.py probe --age 100 + + # 3. Run the study tasks (re-run until it prints 'ALL TASKS DONE'): + python seed_study.py run --age 100 [--budget-seconds 1e9] + + # 4. Summarise results: + python seed_study.py report --age 100 + +Outputs CSV summary to model_output/seed_study/Ma_report.csv. +""" + +import argparse +import json +import marshal +import math +import multiprocessing +import os +import os.path +import pickle +import sys +import time +import types + +import numpy as np +import pygplates as pgp + +import pmagpy.pmag as pmag + +import Optimised_config as cfg + + +STATE_DIR_TEMPLATE = os.path.join('model_output', 'seed_study') + +# Seed-grid sizes to test (requested number of models; for a 'Uniform' search the +# actual number is rounded up to the nearest square). +DEFAULT_SEED_COUNTS = [16, 25, 49, 100, 200] + +# Numbers of best screened seeds to polish (for screen-then-polish strategy). +DEFAULT_TOP_NS = [4, 8, 16, 32] + + +def state_paths(age): + state_dir = STATE_DIR_TEMPLATE + if not os.path.isdir(state_dir): + os.makedirs(state_dir) + return (os.path.join(state_dir, '{0}Ma_progress.json'.format(age)), + os.path.join(state_dir, '{0}Ma_state.pkl'.format(age)), + os.path.join(state_dir, '{0}Ma_tasks.pkl'.format(age)), + os.path.join(state_dir, '{0}Ma_report.csv'.format(age))) + + +def load_json(path, default): + if os.path.exists(path): + with open(path) as f: + return json.load(f) + return default + + +def save_json(path, obj): + tmp = path + '.tmp' + with open(tmp, 'w') as f: + json.dump(obj, f) + os.replace(tmp, path) + + +def load_pickle(path, default=None): + if os.path.exists(path): + with open(path, 'rb') as f: + return pickle.load(f) + return default + + +def save_pickle(path, obj): + tmp = path + '.tmp' + with open(tmp, 'wb') as f: + pickle.dump(obj, f) + os.replace(tmp, path) + + +# --------------------------------------------------------------------------- +# Preparation (chunked, resumable) +# --------------------------------------------------------------------------- + +def cmd_prep(age, budget_seconds): + from optapm import ModelSetup as ms + from no_net_rotation_model import NoNetRotationModel + from optimised_rotation_updater import OptimisedRotationUpdater + from plate_velocity_partitioner import PlateVelocityPartitioner + from continent_fragmentation import ContinentFragmentation + from trench_resolver import TrenchResolver + + t_start = time.time() + progress_path, state_path, _, _ = state_paths(age) + progress = load_json(progress_path, {'nnr_age': None, 'rotfile': False, 'files': False, 'state': False}) + + interval = cfg.interval + ref_rotation_start_age = age + ref_rotation_end_age = age - interval + age_range = range(cfg.actual_end_age + interval, cfg.start_age + interval, interval) + + optimisation_sub_dir = os.path.join(cfg.datadir, cfg.data_model, 'optimisation') + if not os.path.isdir(optimisation_sub_dir): + os.makedirs(optimisation_sub_dir) + + print('Loading topologies...') + sys.stdout.flush() + topology_features = [] + for topology_filename in cfg.topology_filenames: + topology_features.extend(pgp.FeatureCollection(os.path.join(cfg.datadir, topology_filename))) + print(' ...{0:.1f} s'.format(time.time() - t_start)) + sys.stdout.flush() + + # --- Optimised rotation file (created once; re-used on later calls) ---- + # Pass end_age != actual_end_age to re-use existing file (continuation mode). + rot_end_age = age if progress['rotfile'] else cfg.actual_end_age + optimised_rotation_updater = OptimisedRotationUpdater( + cfg.datadir, cfg.original_rotation_filenames, + cfg.start_age, rot_end_age, cfg.actual_end_age, interval, + cfg.get_reference_params, cfg.data_model, cfg.model_name) + rotfile = optimised_rotation_updater.get_optimised_rotation_filename() + if not progress['rotfile']: + progress['rotfile'] = True + save_json(progress_path, progress) + print('Optimised rotation file ready ({0:.1f} s elapsed)'.format(time.time() - t_start)) + sys.stdout.flush() + + # --- No-net-rotation model (chunked) ------------------------------------ + nnr_done_age = progress['nnr_age'] or cfg.actual_end_age + no_net_rotation_model = NoNetRotationModel( + cfg.datadir, cfg.original_rotation_filenames, topology_features, + cfg.start_age, nnr_done_age, cfg.actual_end_age, + cfg.data_model, cfg.model_name) + nnr_rotfile = no_net_rotation_model.get_no_net_rotation_filename() + + while nnr_done_age < age: + if time.time() - t_start > budget_seconds: + progress['nnr_age'] = nnr_done_age + save_json(progress_path, progress) + print('PREP INCOMPLETE (no-net-rotation model at {0}/{1} Ma) - re-run to continue'.format(nnr_done_age, age)) + return + next_age = min(nnr_done_age + 5, age) + no_net_rotation_model.update_no_net_rotation(next_age) + nnr_done_age = next_age + print(' no-net-rotation model updated to {0} Ma ({1:.1f} s elapsed)'.format(nnr_done_age, time.time() - t_start)) + sys.stdout.flush() + progress['nnr_age'] = nnr_done_age + save_json(progress_path, progress) + + # --- Trench + plate velocity files at this age --------------------------- + print('Resolving trenches and plate velocity points...') + sys.stdout.flush() + trench_resolver = TrenchResolver( + cfg.datadir, cfg.original_rotation_filenames, topology_features, cfg.data_model) + tm_file = trench_resolver.get_trench_migration_filename() + trench_resolver.generate_resolved_trenches(ref_rotation_start_age) + + if cfg.plate_velocity_continental_polygons_file: + plate_velocity_plate_features = list( + pgp.FeatureCollection(os.path.join(cfg.datadir, cfg.plate_velocity_continental_polygons_file))) + plate_velocity_features_are_topologies = False + plate_velocity_fragmentation = ContinentFragmentation( + cfg.datadir, cfg.original_rotation_filenames, plate_velocity_plate_features, + cfg.plate_velocity_continental_fragmentation_point_spacing_degrees, + cfg.plate_velocity_continental_fragmentation_area_threshold_steradians, + cfg.plate_velocity_continental_fragmentation_gap_threshold_radians, + age_range) + else: + plate_velocity_plate_features = topology_features + plate_velocity_features_are_topologies = True + plate_velocity_fragmentation = None + + plate_velocity_partitioner = PlateVelocityPartitioner( + cfg.datadir, cfg.original_rotation_filenames, plate_velocity_plate_features, + plate_velocity_features_are_topologies, plate_velocity_fragmentation, + cfg.data_model, cfg.plate_velocity_grid_spacing) + pv_file = plate_velocity_partitioner.get_plate_velocity_filename() + plate_velocity_partitioner.generate_points_and_plate_ids(ref_rotation_start_age) + print(' ...done ({0:.1f} s elapsed)'.format(time.time() - t_start)) + sys.stdout.flush() + + # --- Reference params, seed current interval of optimised rotation ------ + ref_rotation_plate_id, ref_rotation_file = cfg.get_reference_params(ref_rotation_start_age) + if ref_rotation_file == cfg.USE_NNR_REFERENCE_FRAME: + ref_rotation_file = nnr_rotfile + elif ref_rotation_file == cfg.USE_OPTIMISED_REFERENCE_FRAME: + ref_rotation_file = rotfile + + _rotation_model = pgp.RotationModel(os.path.join(cfg.datadir, rotfile)) + plate_rotation_005_rel_000 = _rotation_model.get_rotation(ref_rotation_end_age, 5, fixed_plate_id=0) + plate_rotation_ref_plate_rel_005 = _rotation_model.get_rotation( + ref_rotation_start_age, ref_rotation_plate_id, fixed_plate_id=5) + optimised_rotation_updater.update_optimised_rotation( + plate_rotation_005_rel_000 * plate_rotation_ref_plate_rel_005, + ref_rotation_plate_id, ref_rotation_start_age) + + # --- Component params, data load, starting conditions -------------------- + enable_fz, fz_w, fz_cf, fz_b = cfg.get_fracture_zone_params(ref_rotation_start_age) + enable_nr, nr_w, nr_cf, nr_b = cfg.get_net_rotation_params(ref_rotation_start_age) + enable_tm, tm_w, tm_cf, tm_b = cfg.get_trench_migration_params(ref_rotation_start_age) + enable_hs, hs_w, hs_cf, hs_b = cfg.get_hotspot_trail_params(ref_rotation_start_age) + enable_pv, pv_w, pv_cf, pv_b = cfg.get_plate_velocity_params(ref_rotation_start_age) + + params = [cfg.search_radius, cfg.rotation_uncertainty, cfg.search_type, cfg.models, + cfg.model_stop_condition, cfg.max_iter, + ref_rotation_plate_id, ref_rotation_start_age, ref_rotation_end_age, + cfg.interpolation_resolution, cfg.rotation_age_of_interest, + enable_fz, enable_nr, enable_tm, enable_hs, enable_pv, + fz_w, nr_w, tm_w, hs_w, pv_w, + fz_cf, nr_cf, tm_cf, hs_cf, pv_cf, + fz_b, nr_b, tm_b, hs_b, pv_b, + cfg.no_auto_ref_rot_longitude, cfg.no_auto_ref_rot_latitude, cfg.no_auto_ref_rot_angle, + cfg.auto_calc_ref_pole, cfg.search, + cfg.include_chains, cfg.interpolated_hotspot_trails, cfg.tm_method] + + print('Loading data and computing starting conditions...') + sys.stdout.flush() + data = ms.dataLoader(cfg.datadir, rotfile, ref_rotation_file, + tm_file=tm_file, pv_file=pv_file, nnr_rotfile=nnr_rotfile, + ridge_file=cfg.ridge_file if enable_fz else None, + isochron_file=cfg.isochron_file if enable_fz else None, + isocob_file=cfg.isocob_file if enable_fz else None, + hst_file=cfg.hst_file, hs_file=cfg.hs_file, + interpolated_hotspots=cfg.interpolated_hotspots) + sc = ms.modelStartConditions(params, data, False) + + (x, opt_n, N, lb, ub, model_stop_condition, max_iter, + rotation_file, _, _, _, + Lats, Lons, spreading_directions, spreading_asymmetries, seafloor_ages, PID, CPID, + data_array, weights_array, cost_func_array, bounds_array, + trench_migration_file, plate_velocity_file, no_net_rotation_file, reformArray, trail_data, + start_seeds, rotation_age_of_interest_age, data_array_labels_short, + ref_rot_longitude, ref_rot_latitude, ref_rot_angle, seed_lons, seed_lats) = sc[:35] + + state = { + 'interval': cfg.interval, + 'rotation_file': rotation_file, + 'no_net_rotation_file': no_net_rotation_file, + 'ref_rotation_start_age': ref_rotation_start_age, + 'ref_rotation_end_age': ref_rotation_end_age, + 'ref_rotation_plate_id': ref_rotation_plate_id, + 'Lats': Lats, 'Lons': Lons, + 'spreading_directions': spreading_directions, + 'spreading_asymmetries': spreading_asymmetries, + 'seafloor_ages': seafloor_ages, 'PID': PID, 'CPID': CPID, + 'data_array': data_array, 'weights_array': weights_array, + 'cost_func_code_strings': [marshal.dumps(cf.__code__) for cf in cost_func_array], + 'bounds_array': bounds_array, + 'trench_migration_file': trench_migration_file, + 'plate_velocity_file': plate_velocity_file, + 'reformArray': reformArray, 'trail_data': trail_data, + 'use_trail_age_uncertainty': cfg.use_trail_age_uncertainty, + 'trail_age_uncertainty_ellipse': cfg.trail_age_uncertainty_ellipse, + 'tm_method': cfg.tm_method, + 'x': [list(map(float, xi)) for xi in x], + 'opt_n': opt_n, 'lb': list(lb), 'ub': list(ub), + 'model_stop_condition': model_stop_condition, 'max_iter': max_iter, + 'ref_rot_longitude': float(ref_rot_longitude), + 'ref_rot_latitude': float(ref_rot_latitude), + 'ref_rot_angle': float(ref_rot_angle), + 'search_radius': cfg.search_radius, + } + save_pickle(state_path, state) + progress['files'] = True + progress['state'] = True + save_json(progress_path, progress) + print('PREP DONE ({0:.1f} s total, {1} seeds)'.format(time.time() - t_start, len(x))) + + +# --------------------------------------------------------------------------- +# Worker processes +# --------------------------------------------------------------------------- + +_g_obj_f = None +_g_nlopt_args = None + + +def _init_worker(state): + global _g_obj_f, _g_nlopt_args + from objective_function import ObjectiveFunction + cost_func_array = [types.FunctionType(marshal.loads(code), globals(), 'cost_func') + for code in state['cost_func_code_strings']] + _g_obj_f = ObjectiveFunction( + state['interval'], state['rotation_file'], state['no_net_rotation_file'], + state['ref_rotation_start_age'], state['Lats'], state['Lons'], + state['spreading_directions'], state['spreading_asymmetries'], state['seafloor_ages'], + state['PID'], state['CPID'], state['data_array'], state['weights_array'], + cost_func_array, state['bounds_array'], + state['trench_migration_file'], state['plate_velocity_file'], + state['ref_rotation_end_age'], state['ref_rotation_plate_id'], + state['reformArray'], state['trail_data'], + state['use_trail_age_uncertainty'], state['trail_age_uncertainty_ellipse'], + state['tm_method']) + _g_nlopt_args = (state['opt_n'], state['lb'], state['ub'], + state['model_stop_condition'], state['max_iter']) + + +def _run_task(task): + kind = task['kind'] + seed = task['seed'] + t0 = time.time() + if kind == 'screen': + cost = _g_obj_f(seed, None) + return task['id'], {'cost': float(cost), 'time': time.time() - t0} + else: # full optimisation + import nlopt + opt_n, lb, ub, model_stop_condition, max_iter = _g_nlopt_args + opt = nlopt.opt(nlopt.LN_COBYLA, opt_n) + opt.set_min_objective(_g_obj_f) + opt.set_lower_bounds(lb) + opt.set_upper_bounds(ub) + if model_stop_condition != 'threshold': + opt.set_maxeval(max_iter) + else: + opt.set_ftol_rel(1e-6) + opt.set_xtol_rel(1e-8) + # Safety cap: COBYLA occasionally fails to converge (cycles for thousands of + # evaluations). The vast majority of optimisations converge in < 200 evaluations, + # so 300 is a generous cap that only affects pathological cases. + opt.set_maxeval(300) + xopt = opt.optimize(list(seed)) + return task['id'], {'xopt': [float(v) for v in xopt], + 'cost': float(opt.last_optimum_value()), + 'nevals': opt.get_numevals(), + 'time': time.time() - t0} + + +# --------------------------------------------------------------------------- +# Probe +# --------------------------------------------------------------------------- + +def cmd_probe(age): + _, state_path, _, _ = state_paths(age) + state = load_pickle(state_path) + if state is None: + print('No state found - run prep first.') + return + _init_worker(state) + x = state['x'] + # Time single evaluations. + times = [] + for seed in x[:3]: + t0 = time.time() + cost = _g_obj_f(seed, None) + times.append(time.time() - t0) + print('eval at ({0:.1f}, {1:.1f}, {2:.2f}) -> cost {3:.4f} in {4:.3f} s'.format( + seed[0], seed[1], seed[2], cost, times[-1])) + # Time one full optimisation. + _, res = _run_task({'kind': 'full', 'seed': x[0], 'id': 'probe'}) + print('full COBYLA: cost {0:.4f}, {1} evals, {2:.1f} s'.format(res['cost'], res['nevals'], res['time'])) + print('mean eval time: {0:.3f} s'.format(np.mean(times))) + + +# --------------------------------------------------------------------------- +# Run (chunked, resumable) +# --------------------------------------------------------------------------- + +def generate_uniform_seeds(num_models, search_radius, ref_rot_longitude, ref_rot_latitude, ref_rot_angle): + """Uniform seed grid exactly as ModelSetup.modelStartConditions (search_type='Uniform').""" + num_points_float = num_models * 1.0 / (0.5 * (1 - math.cos(math.radians(search_radius)))) + n = int(math.ceil(math.sqrt(num_points_float))) + seeds = [] + for lon_index in range(n): + for lat_index in range(n): + theta = 2 * np.pi * ((lon_index + 0.5) / float(n)) + phi = np.arccos(2 * ((lat_index + 0.5) / float(n)) - 1.0) + xyz = (np.cos(theta) * np.sin(phi), np.sin(theta) * np.sin(phi), np.cos(phi)) + point = pgp.convert_point_on_sphere_to_lat_lon_point(xyz) + lat, lon = point.get_latitude(), point.get_longitude() + if lat > 90 - search_radius: + rlon, rlat = pmag.dodirot(lon, lat, ref_rot_longitude, ref_rot_latitude) + seeds.append([float(rlon), float(rlat), float(ref_rot_angle)]) + return seeds + + +def build_tasks(state, seed_counts, top_ns): + """Build the full task list (screen + full multistart + reduced grids).""" + tasks = {} + x = state['x'] + for i, seed in enumerate(x): + tasks['screen_{0}'.format(i)] = {'kind': 'screen', 'seed': seed, 'id': 'screen_{0}'.format(i)} + tasks['full_{0}'.format(i)] = {'kind': 'full', 'seed': seed, 'id': 'full_{0}'.format(i)} + for n in seed_counts: + seeds_n = generate_uniform_seeds(n, state['search_radius'], + state['ref_rot_longitude'], state['ref_rot_latitude'], + state['ref_rot_angle']) + for i, seed in enumerate(seeds_n): + tid = 'grid{0}_{1}'.format(n, i) + tasks[tid] = {'kind': 'full', 'seed': seed, 'id': tid} + # Note: 'polish' tasks for screen-then-polish are just the 'full_' tasks of the + # top-N screened seeds, so they don't need separate optimisations - the report + # simply combines screening costs with the relevant full results. + return tasks + + +def cmd_run(age, budget_seconds, processes, seed_counts, top_ns): + t_start = time.time() + progress_path, state_path, tasks_path, _ = state_paths(age) + state = load_pickle(state_path) + if state is None: + print('No state found - run prep first.') + return + + task_store = load_pickle(tasks_path) + if task_store is None: + tasks = build_tasks(state, seed_counts, top_ns) + task_store = {'tasks': tasks, 'done': {}, 'seed_counts': seed_counts, 'top_ns': top_ns} + save_pickle(tasks_path, task_store) + + tasks = task_store['tasks'] + done = task_store['done'] + pending = [tid for tid in tasks if tid not in done] + # Run screening tasks first (cheap, and needed for screen-then-polish). + pending.sort(key=lambda tid: (0 if tasks[tid]['kind'] == 'screen' else 1, tid)) + print('{0} tasks pending, {1} done'.format(len(pending), len(done))) + if not pending: + print('ALL TASKS DONE') + return + + pool = multiprocessing.Pool(processes, initializer=_init_worker, initargs=(state,)) + n_completed = 0 + try: + result_iter = pool.imap_unordered(_run_task, (tasks[tid] for tid in pending)) + for tid, res in result_iter: + done[tid] = res + n_completed += 1 + if n_completed % 20 == 0: + save_pickle(tasks_path, task_store) + if time.time() - t_start > budget_seconds: + break + finally: + pool.terminate() + pool.join() + save_pickle(tasks_path, task_store) + + remaining = len(tasks) - len(done) + print('completed {0} tasks this call ({1:.1f} s); {2} remaining'.format( + n_completed, time.time() - t_start, remaining)) + if remaining == 0: + print('ALL TASKS DONE') + + +# --------------------------------------------------------------------------- +# Report +# --------------------------------------------------------------------------- + +def cmd_report(age): + _, state_path, tasks_path, report_path = state_paths(age) + state = load_pickle(state_path) + task_store = load_pickle(tasks_path) + if state is None or task_store is None: + print('No state/tasks found - run prep and run first.') + return + tasks, done = task_store['tasks'], task_store['done'] + seed_counts, top_ns = task_store['seed_counts'], task_store['top_ns'] + x = state['x'] + n_seeds = len(x) + + def pole_dist_deg(a, b): + import geoTools + lat1, lon1 = geoTools.checkLatLon(a[1], a[0]) + lat2, lon2 = geoTools.checkLatLon(b[1], b[0]) + p1 = pgp.PointOnSphere(lat1, lon1) + p2 = pgp.PointOnSphere(lat2, lon2) + return math.degrees(pgp.GeometryOnSphere.distance(p1, p2)) + + rows = [] + + # Ground truth: full multistart. + full_results = [done.get('full_{0}'.format(i)) for i in range(n_seeds)] + have_full = [r for r in full_results if r] + if not have_full: + print('No full multistart results yet.') + return + gt = min(have_full, key=lambda r: r['cost']) + gt_cost, gt_x = gt['cost'], gt['xopt'] + total_full_time = sum(r['time'] for r in have_full) + total_full_evals = sum(r['nevals'] for r in have_full) + mean_opt_time = total_full_time / len(have_full) + print('Ground truth ({0}/{1} full optimisations): cost {2:.6f} at ' + 'lon {3:.2f} lat {4:.2f} ang {5:.3f}'.format( + len(have_full), n_seeds, gt_cost, gt_x[0], gt_x[1], gt_x[2])) + print(' mean {0:.0f} evals and {1:.1f} s per optimisation; total {2:.0f} s CPU'.format( + total_full_evals / len(have_full), mean_opt_time, total_full_time)) + rows.append(['full_multistart', n_seeds, len(have_full), '{0:.6f}'.format(gt_cost), + '{0:.3f}'.format(gt_x[0]), '{0:.3f}'.format(gt_x[1]), '{0:.4f}'.format(gt_x[2]), + '0.0', '0.0', '{0:.1f}'.format(total_full_time)]) + + # Reduced grids. + for n in seed_counts: + grid_results = [done[tid] for tid in done if tid.startswith('grid{0}_'.format(n))] + n_grid_tasks = sum(1 for tid in tasks if tid.startswith('grid{0}_'.format(n))) + if not grid_results or len(grid_results) < n_grid_tasks: + print('grid {0}: incomplete ({1}/{2})'.format(n, len(grid_results), n_grid_tasks)) + continue + best = min(grid_results, key=lambda r: r['cost']) + t_total = sum(r['time'] for r in grid_results) + cost_diff_pct = 100.0 * (best['cost'] - gt_cost) / abs(gt_cost) + pdist = pole_dist_deg(best['xopt'], gt_x) + print('grid {0:4d} ({1:3d} seeds): best cost {2:.6f} (+{3:.3f}%), pole dist {4:6.2f} deg, ' + 'ang diff {5:+.3f}, CPU {6:.0f} s'.format( + n, len(grid_results), best['cost'], cost_diff_pct, pdist, + best['xopt'][2] - gt_x[2], t_total)) + rows.append(['uniform_grid_{0}'.format(n), n, len(grid_results), '{0:.6f}'.format(best['cost']), + '{0:.3f}'.format(best['xopt'][0]), '{0:.3f}'.format(best['xopt'][1]), + '{0:.4f}'.format(best['xopt'][2]), + '{0:.4f}'.format(cost_diff_pct), '{0:.2f}'.format(pdist), '{0:.1f}'.format(t_total)]) + + # Screen-then-polish. + screen_results = [done.get('screen_{0}'.format(i)) for i in range(n_seeds)] + if all(screen_results): + screen_costs = np.array([r['cost'] for r in screen_results]) + t_screen = sum(r['time'] for r in screen_results) + order = np.argsort(screen_costs) + for top_n in top_ns: + top_idx = order[:top_n] + polish = [done.get('full_{0}'.format(i)) for i in top_idx] + if not all(polish): + print('screen+polish top {0}: polish results incomplete'.format(top_n)) + continue + best = min(polish, key=lambda r: r['cost']) + t_total = t_screen + sum(r['time'] for r in polish) + cost_diff_pct = 100.0 * (best['cost'] - gt_cost) / abs(gt_cost) + pdist = pole_dist_deg(best['xopt'], gt_x) + print('screen all {0} seeds + polish top {1:3d}: best cost {2:.6f} (+{3:.3f}%), ' + 'pole dist {4:6.2f} deg, CPU {5:.0f} s ({6:.1f}x speedup vs full)'.format( + n_seeds, top_n, best['cost'], cost_diff_pct, pdist, t_total, + total_full_time / t_total)) + rows.append(['screen_polish_top{0}'.format(top_n), n_seeds, top_n, + '{0:.6f}'.format(best['cost']), + '{0:.3f}'.format(best['xopt'][0]), '{0:.3f}'.format(best['xopt'][1]), + '{0:.4f}'.format(best['xopt'][2]), + '{0:.4f}'.format(cost_diff_pct), '{0:.2f}'.format(pdist), '{0:.1f}'.format(t_total)]) + else: + print('screening incomplete: {0}/{1}'.format(sum(1 for r in screen_results if r), n_seeds)) + + with open(report_path, 'w') as f: + f.write('strategy,seeds_requested,optimisations_run,best_cost,lon,lat,angle,cost_diff_pct,pole_dist_deg,cpu_time_s\n') + for row in rows: + f.write(','.join(str(v) for v in row) + '\n') + print('Wrote {0}'.format(report_path)) + + +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser(description='optAPM seed-count study (resumable)') + parser.add_argument('command', choices=['prep', 'probe', 'run', 'report']) + parser.add_argument('--age', type=int, required=True) + parser.add_argument('--budget-seconds', type=float, default=1e9, + help='Approximate wall-time budget for this call (prep/run are resumable)') + parser.add_argument('--processes', type=int, default=multiprocessing.cpu_count()) + parser.add_argument('--seed-counts', type=int, nargs='+', default=DEFAULT_SEED_COUNTS) + parser.add_argument('--top-ns', type=int, nargs='+', default=DEFAULT_TOP_NS) + args = parser.parse_args() + + if args.age % cfg.interval != 0: + raise SystemExit('age must be a multiple of the interval ({0})'.format(cfg.interval)) + + if args.command == 'prep': + cmd_prep(args.age, args.budget_seconds) + elif args.command == 'probe': + cmd_probe(args.age) + elif args.command == 'run': + cmd_run(args.age, args.budget_seconds, args.processes, args.seed_counts, args.top_ns) + elif args.command == 'report': + cmd_report(args.age) + + +if __name__ == '__main__': + main() From 8ecd55745cb1315bd41482130ad7173acb1c2ad9 Mon Sep 17 00:00:00 2001 From: Dietmar Muller Date: Thu, 11 Jun 2026 23:54:13 +0000 Subject: [PATCH 3/6] Add run variants for uncertainty quantification ('best' and 'nr_max') The optimisation is dominated by the net rotation (NR) minimisation (subduction zone migration and continent speeds are not independent of it - see Muller et al. 2022, Solid Earth), so a defensible uncertainty envelope for the optimised reference frame consists of end-members in the NR bounds rather than perturbations of the component weights: 1. The no-net-rotation reference frame (zero-NR end-member) - already produced by every run. 2. The 'best' run (default) - NR bounds (0.08, 0.20) deg/Myr as before. 3. The 'nr_max' run - NR upper bound relaxed to 0.30 deg/Myr, the top of the range permitted by asthenospheric shear / seismic anisotropy constraints (Conrad & Behn 2010; Becker 2006). The variant is selected via 'run_variant' in Optimised_config.py or the OPTAPM_VARIANT environment variable, and (for variants other than 'best') is appended to the model name so each variant writes its own output files. --- Optimised_config.py | 42 +++++++++++++++++++++++++++++++++++++++++- README.md | 17 +++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/Optimised_config.py b/Optimised_config.py index 6a3906b..fc26585 100644 --- a/Optimised_config.py +++ b/Optimised_config.py @@ -50,6 +50,45 @@ model_name = "run1" +# +# Run variant (for uncertainty quantification). +# +# The optimisation is dominated by the net rotation (NR) minimisation: subduction zone +# migration largely depends on the net rotation optimisation (it is not really an +# independent parameter), and the same holds for limiting the speed of continents +# (see Muller et al. 2022, Solid Earth, doi:10.5194/se-13-1127-2022). +# So a defensible uncertainty envelope for the optimised reference frame consists of +# end-members in the net rotation bounds rather than perturbations of the weights: +# +# 1. A no-net-rotation (NNR) reference frame - the zero-NR end-member. This is already +# produced by every run as "no_net_rotation_model_.rot". +# 2. The 'best' run - the default NR bounds of (0.08, 0.20) deg/Myr, ie, non-zero +# (as suggested by mantle flow models, eg, Becker 2006) but below the preferred +# geodynamic upper limit (Conrad & Behn 2010). +# 3. The 'nr_max' run - the maximum-NR end-member, with the NR upper bound relaxed to +# 0.30 deg/Myr: the top of the range permitted by asthenospheric shear / seismic +# anisotropy constraints (Conrad & Behn 2010 derive NR < 0.26 deg/Myr for an +# asthenosphere 10x less viscous than the upper mantle; anisotropy proxies allow +# 0.2-0.3 deg/Myr). Beyond this, net rotations approach the Pacific hotspot (HS3) +# frame (0.33-0.44 deg/Myr), which is widely considered geodynamically implausible. +# +# Select the variant here (or via the OPTAPM_VARIANT environment variable, eg, +# "OPTAPM_VARIANT=nr_max mpirun -np 16 python Optimised_APM.py"). +# The variant name (other than 'best') is appended to the model name, so each variant +# writes its own set of output files. +# +run_variant_nr_bounds = { + 'best': (0.08, 0.20), # deg/Myr + 'nr_max': (0.08, 0.30), # deg/Myr +} +run_variant = os.environ.get('OPTAPM_VARIANT', 'best') +if run_variant not in run_variant_nr_bounds: + raise RuntimeError('Unknown run variant {0!r} - choose one of {1}'.format( + run_variant, sorted(run_variant_nr_bounds.keys()))) +if run_variant != 'best': + model_name = model_name + '_' + run_variant + + # Start age. if data_model == 'Zahirovic_etal_2022_GDJ': start_age = 410 @@ -299,7 +338,8 @@ def cost_function(PTLong1, PTLat1, PTangle1, SPLong, SPLat, SPangle, SPLong_NNR, data_model == 'Zahirovic_etal_2022_GDJ' or '1.8ga' in data_model.lower() or data_model == "Global_2000-540"): - nr_bounds = (0.08, 0.20) + # Net rotation bounds depend on the run variant (see 'run_variant' above). + nr_bounds = run_variant_nr_bounds[run_variant] if age <= 80: return True, 1.0, cost_function, nr_bounds # Weight is always 1.0 for 0-80Ma else: diff --git a/README.md b/README.md index fe95837..c7df0e5 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,23 @@ To disable screening (and fully optimise every seed, as in the original workflow The workflow also caps each NLopt (COBYLA) optimisation at `nlopt_max_eval_safety` objective evaluations (default 1000). COBYLA occasionally fails to converge and can otherwise cycle for thousands of evaluations, stalling an entire MPI rank (and hence the whole timestep, since all ranks must finish before the next timestep begins). +### Uncertainty quantification (run variants) + +The optimisation is dominated by the net rotation (NR) minimisation: subduction zone migration largely depends on the net rotation optimisation (it is not an independent parameter), and the same holds for limiting the speed of continents (see [Muller et al. 2022, Solid Earth](https://doi.org/10.5194/se-13-1127-2022)). Perturbing the component weights only slightly changes the outcome, so a defensible uncertainty envelope for the optimised reference frame instead consists of end-members in the net rotation bounds: + +1. **No-net-rotation (NNR) end-member** - zero net rotation. This rotation file is already produced by every run as "no_net_rotation_model_.rot". +2. **Best run** (`run_variant = 'best'`, the default) - NR bounded to (0.08, 0.20) deg/Myr: non-zero (as suggested by mantle flow models, eg, Becker 2006) but below the preferred geodynamic upper limit (Conrad & Behn 2010). +3. **Maximum-NR end-member** (`run_variant = 'nr_max'`) - the NR upper bound relaxed to 0.30 deg/Myr, the top of the range permitted by asthenospheric shear / seismic anisotropy constraints (Conrad & Behn 2010 derive NR < 0.26 deg/Myr for an asthenosphere 10x less viscous than the upper mantle; anisotropy proxies allow 0.2-0.3 deg/Myr). Larger values approach the Pacific hotspot (HS3) frame rates (0.33-0.44 deg/Myr), which are widely considered geodynamically implausible. + +To produce the envelope, run the workflow twice (the `nr_max` variant appends its name to the model name, so the two runs write separate output files): + +``` + mpirun -np python Optimised_APM.py + OPTAPM_VARIANT=nr_max mpirun -np python Optimised_APM.py +``` + +The variant NR bounds are defined in `run_variant_nr_bounds` in "Optimised_config.py" (where further variants can be added). + ### Quick test runs The age range, seed count, screening parameters and parallelisation can be temporarily overridden via environment variables without editing "Optimised_config.py" - useful for quick smoke tests: From 753f4831b84b82def9e3314819d1d4102afdc32e Mon Sep 17 00:00:00 2001 From: Dietmar Muller Date: Fri, 12 Jun 2026 00:18:45 +0000 Subject: [PATCH 4/6] Prefer slow trench rollback in the trench migration component (new default) The original trench migration component drives trench-normal migration velocities toward zero. Since zero net rotation also implies near-zero trench migration, that constraint is largely redundant with the net rotation minimisation. It is also at odds with observations: most trenches roll back toward the ocean basin behind them rather than advancing toward the overriding plate - across reference frames, 62-78% of trench segments retreat, with mean trench-normal velocities of +1.3-1.5 cm/yr and medians of +0.9-1.3 cm/yr (Schellart et al. 2008, Earth-Science Reviews; Williams et al. 2015, EPSL). The new default scheme ('trench_migration_scheme = "rollback"') drives per-trench orthogonal velocities toward a target of +10 mm/yr retreat and tightens the bounds on the mean orthogonal velocity from (-30, 30) to (0, 20) mm/yr - the global mean must be retreating, but at less than 2 cm/yr. Individual trenches can still advance; only the mean is required to retreat. This makes trench kinematics an independent constraint that competes with the net rotation minimisation (instead of duplicating it). Set 'trench_migration_scheme = "minimise"' (or OPTAPM_TM_SCHEME=minimise) to restore the original behaviour. --- Optimised_config.py | 65 +++++++++++++++++++++++++++++++++++++++++++++ README.md | 6 +++++ 2 files changed, 71 insertions(+) diff --git a/Optimised_config.py b/Optimised_config.py index fc26585..c34acc2 100644 --- a/Optimised_config.py +++ b/Optimised_config.py @@ -347,7 +347,68 @@ def cost_function(PTLong1, PTLat1, PTangle1, SPLong, SPLat, SPangle, SPLong_NNR, else: return True, 1.0, cost_function, None +# +# Trench migration scheme. +# +# 'minimise' - the original behaviour: drive trench-normal migration velocities toward zero +# (mean absolute orthogonal velocity + standard deviation), with loose bounds +# of (-30, 30) mm/yr on the mean orthogonal velocity. +# +# 'rollback' - prefer slow trench *retreat* (rollback). Observations show that most trenches +# roll back toward the ocean basin behind them rather than advancing toward the +# overriding plate: across reference frames, 62-78% of trench segments retreat, +# with mean trench-normal velocities of +1.3-1.5 cm/yr and medians of +# +0.9-1.3 cm/yr (Schellart et al. 2008, Earth-Science Reviews, +# doi:10.1016/j.earscirev.2008.01.005; see also Williams et al. 2015, EPSL, +# doi:10.1016/j.epsl.2015.02.026). This scheme drives the per-trench orthogonal +# velocities toward a target of +10 mm/yr (positive = retreat; the target value +# is hardcoded inside the cost function below because cost functions are +# serialised by code object only and cannot capture config variables), and +# tightens the bounds on the *mean* orthogonal velocity to (0, 20) mm/yr, +# ie, the global mean must be retreating but at less than 2 cm/yr. +# Individual trenches can still advance (eg, Izu-Bonin-Mariana-style segments); +# only the global mean is required to be retreating. +# +# Because zero net rotation also implies near-zero trench migration, the original +# 'minimise' scheme is largely redundant with net rotation minimisation. The +# 'rollback' scheme makes trench kinematics an independent constraint that +# competes with the net rotation minimisation. +# +trench_migration_scheme = 'rollback' + def get_trench_migration_params(age): + if trench_migration_scheme == 'rollback': + # Cost function - see "objective_function.py" for definition of function arguments... + def cost_function(trench_vel, trench_obl, tm_vel_orth, tm_mean_vel_orth, tm_mean_abs_vel_orth): + # NOTE: Import any modules used in this function here + # (since this function's code might be serialised over the network to remote nodes). + # + # NOTE: Any constants must be hardcoded here (not read from the config) since only + # this function's *code* is serialised (closures/globals are not preserved). + import numpy as np + + # Target trench-normal migration velocity in mm/yr (positive = retreat/rollback). + # The median trench-normal retreat is +0.9-1.3 cm/yr across reference frames + # (Schellart et al. 2008), so we target +10 mm/yr (= +1.0 cm/yr). + tm_target = 10.0 + + # Mean absolute deviation of per-trench orthogonal velocities from the rollback + # target, plus the standard deviation (to penalise large spread), with the same + # overall scaling as the 'minimise' scheme so the component weights are comparable. + return (np.mean(np.abs(tm_vel_orth - tm_target)) + np.std(tm_vel_orth)) * 3 + + # Bounds on the *mean* orthogonal velocity (mm/yr): must be retreating, but slower + # than 2 cm/yr. (Note: Use units of mm/yr - same as km/Myr.) + tm_bounds = [0, 20] + + if age <= 80: + return True, 1.0, cost_function, tm_bounds # Weight is always 1.0 for 0-80Ma + else: + # NOTE: These are inverse weights (ie, the constraint costs are *multiplied* by "1.0 / weight"). + return True, 2.0, cost_function, tm_bounds # 2.0 gives a *multiplicative* weight of 0.5 + + # trench_migration_scheme == 'minimise' (the original scheme) ... + # Cost function - see "objective_function.py" for definition of function arguments... def cost_function(trench_vel, trench_obl, tm_vel_orth, tm_mean_vel_orth, tm_mean_abs_vel_orth): # NOTE: Import any modules used in this function here @@ -655,6 +716,10 @@ def get_reference_params(age): seed_screen_top_n = None if 'OPTAPM_SCREEN_UNIFORM_N' in os.environ: seed_screen_uniform_n = int(os.environ['OPTAPM_SCREEN_UNIFORM_N']) +if 'OPTAPM_TM_SCHEME' in os.environ: + trench_migration_scheme = os.environ['OPTAPM_TM_SCHEME'] + if trench_migration_scheme not in ('minimise', 'rollback'): + raise RuntimeError("OPTAPM_TM_SCHEME must be 'minimise' or 'rollback'") if os.environ.get('OPTAPM_SERIAL') == '1': use_parallel = None diff --git a/README.md b/README.md index c7df0e5..11eb74b 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,12 @@ To disable screening (and fully optimise every seed, as in the original workflow The workflow also caps each NLopt (COBYLA) optimisation at `nlopt_max_eval_safety` objective evaluations (default 1000). COBYLA occasionally fails to converge and can otherwise cycle for thousands of evaluations, stalling an entire MPI rank (and hence the whole timestep, since all ranks must finish before the next timestep begins). +### Trench migration scheme + +The original trench migration component drives trench-normal migration velocities toward zero. However, since zero net rotation also implies near-zero trench migration, that constraint is largely redundant with the net rotation minimisation. It is also at odds with observations: most trenches roll back toward the ocean basin behind them rather than advancing toward the overriding plate - across reference frames, 62-78% of trench segments retreat, with mean trench-normal velocities of +1.3-1.5 cm/yr and medians of +0.9-1.3 cm/yr (Schellart et al. 2008, Earth-Science Reviews; see also Williams et al. 2015, EPSL). + +The default scheme is therefore now `trench_migration_scheme = 'rollback'` (in "Optimised_config.py"): per-trench orthogonal velocities are driven toward a target of +10 mm/yr retreat, and the bounds on the *mean* orthogonal velocity are tightened to (0, 20) mm/yr - the global mean must be retreating, but at less than 2 cm/yr. Individual trenches can still advance; only the mean is required to retreat. This makes trench kinematics an independent constraint that competes with the net rotation minimisation (instead of duplicating it). Set `trench_migration_scheme = 'minimise'` (or `OPTAPM_TM_SCHEME=minimise`) to restore the original behaviour. + ### Uncertainty quantification (run variants) The optimisation is dominated by the net rotation (NR) minimisation: subduction zone migration largely depends on the net rotation optimisation (it is not an independent parameter), and the same holds for limiting the speed of continents (see [Muller et al. 2022, Solid Earth](https://doi.org/10.5194/se-13-1127-2022)). Perturbing the component weights only slightly changes the outcome, so a defensible uncertainty envelope for the optimised reference frame instead consists of end-members in the net rotation bounds: From 6ad8c4626fb72d5ac099cc501075111af38eb713 Mon Sep 17 00:00:00 2001 From: Dietmar Muller Date: Fri, 12 Jun 2026 00:36:57 +0000 Subject: [PATCH 5/6] Strengthen the 'nr_max' uncertainty variant: also halve the net rotation weight Relaxing the net rotation bounds alone is not sufficient to produce a directional end-member: the bounds are hard penalty walls that only affect timesteps where the optimum presses against the 0.20 deg/Myr ceiling (mostly in deep time), whereas inside the walls the solution is set by the smooth trade-off between the NR cost and the other components. Halving the NR weight (inverse weight 2.0) shifts that trade-off at every timestep, letting the trench rollback and plate velocity constraints pull harder against the net rotation minimisation. Single-timestep tests (5 Ma) confirm: bounds-only relaxation merely selected a different near-degenerate minimum (NR 0.123 deg/Myr), whereas bounds + halved weight produced the intended systematically stronger-rollback, higher-NR end-member (NR 0.177 deg/Myr, trench retreat strengthening from +1.8 to +3.2 mm/yr, 68% of segments retreating). --- Optimised_config.py | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/Optimised_config.py b/Optimised_config.py index c34acc2..cf9f604 100644 --- a/Optimised_config.py +++ b/Optimised_config.py @@ -72,6 +72,17 @@ # 0.2-0.3 deg/Myr). Beyond this, net rotations approach the Pacific hotspot (HS3) # frame (0.33-0.44 deg/Myr), which is widely considered geodynamically implausible. # +# The 'nr_max' variant also *halves the net rotation weight* (inverse weight 2.0). +# Relaxing the bounds alone is not sufficient to produce a directional end-member: +# the bounds are hard penalty walls that only affect timesteps where the optimum +# presses against the 0.20 deg/Myr ceiling (mostly in deep time), whereas inside the +# walls the solution is set by the smooth trade-off between the net rotation cost and +# the other components. Halving the NR weight shifts that trade-off at *every* +# timestep, letting the trench rollback and plate velocity constraints pull harder +# against the net rotation minimisation (single-timestep tests confirm this produces +# a systematically stronger-rollback / higher-NR solution, whereas relaxing the +# bounds alone merely selects a different near-degenerate minimum). +# # Select the variant here (or via the OPTAPM_VARIANT environment variable, eg, # "OPTAPM_VARIANT=nr_max mpirun -np 16 python Optimised_APM.py"). # The variant name (other than 'best') is appended to the model name, so each variant @@ -81,6 +92,11 @@ 'best': (0.08, 0.20), # deg/Myr 'nr_max': (0.08, 0.30), # deg/Myr } +# Inverse weights (the NR cost is *multiplied* by "1.0 / weight"). +run_variant_nr_weight = { + 'best': 1.0, # full NR weight + 'nr_max': 2.0, # half NR weight +} run_variant = os.environ.get('OPTAPM_VARIANT', 'best') if run_variant not in run_variant_nr_bounds: raise RuntimeError('Unknown run variant {0!r} - choose one of {1}'.format( @@ -338,12 +354,13 @@ def cost_function(PTLong1, PTLat1, PTangle1, SPLong, SPLat, SPangle, SPLong_NNR, data_model == 'Zahirovic_etal_2022_GDJ' or '1.8ga' in data_model.lower() or data_model == "Global_2000-540"): - # Net rotation bounds depend on the run variant (see 'run_variant' above). + # Net rotation bounds and weight depend on the run variant (see 'run_variant' above). nr_bounds = run_variant_nr_bounds[run_variant] + nr_weight = run_variant_nr_weight[run_variant] if age <= 80: - return True, 1.0, cost_function, nr_bounds # Weight is always 1.0 for 0-80Ma + return True, nr_weight, cost_function, nr_bounds else: - return True, 1.0, cost_function, nr_bounds # 1.0 gives a *multiplicative* weight of 1.0 + return True, nr_weight, cost_function, nr_bounds # NOTE: Inverse weight (cost is multiplied by "1.0 / weight"). else: return True, 1.0, cost_function, None @@ -632,6 +649,13 @@ def get_reference_params(age): nlopt_max_eval_safety = 1000 +# Whether to generate post-run diagnostics at the end of the optimisation run: +# per-timestep net rotation (median +/- MAD) and trench-normal migration (mean +/- MAD) +# statistics (CSV) and plots (PNG) written to 'model_output/' - see "model_diagnostics.py". +# The diagnostics can also be (re)generated standalone: "python model_diagnostics.py". +generate_diagnostics = True + + search = "Initial" # If True then temporarily expand search radius to 180 whenever the reference plate changes. # Normally the reference plate stays constant at Africa (701), but does switch to 101 for the 1Ga model. From f2bf2dcbace221b99f736a146d198cf3ee868869 Mon Sep 17 00:00:00 2001 From: Dietmar Muller Date: Fri, 12 Jun 2026 00:36:57 +0000 Subject: [PATCH 6/6] Add post-run model diagnostics (net rotation and trench migration through time) At the end of every run (config flag 'generate_diagnostics', default True) the workflow now computes per-timestep summary statistics of the optimised model and writes them to 'model_output/': * _diagnostics.csv - per timestep: median and MAD of the net rotation rate (deg/Myr, sampled at 1 Myr sub-steps within each timestep), and mean/median/MAD of the trench-normal migration velocities (mm/yr, positive = retreat) across all subduction zone segments, plus the percentage of segments retreating. * _net_rotation.png - NR (median +/- MAD) through time, with the 0.20 deg/Myr preferred upper bound and the 0.26 deg/Myr Conrad & Behn (2010) limit marked. * _trench_migration.png - trench-normal migration (mean +/- MAD) through time, with the observed present-day mean retreat range (+0.9 to +1.5 cm/yr, Schellart et al. 2008) marked. The same statistics/plots are also generated for the no-net-rotation model (outputs '_NNR_*') - showing subduction zone kinematics in an NNR world (the zero-NR end-member of the uncertainty envelope), with the NNR net rotation plot doubling as a sanity check (~zero at all times). Diagnostics can be (re)generated standalone on an existing optimised model via 'python model_diagnostics.py'. Also added a 'References for parameter choices' section to the README and expanded the justifications (with references and validation tables) for all recently introduced parameters: seed screening counts, COBYLA evaluation cap, net rotation bounds and weights per run variant, and the trench rollback target and bounds. --- Optimised_APM.py | 16 ++- README.md | 53 ++++++++- model_diagnostics.py | 278 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 345 insertions(+), 2 deletions(-) create mode 100644 model_diagnostics.py diff --git a/Optimised_APM.py b/Optimised_APM.py index c0bcab8..21f3c91 100644 --- a/Optimised_APM.py +++ b/Optimised_APM.py @@ -808,7 +808,21 @@ def select_polish_seeds(seeds, costs, top_n, uniform_n): # Save the final optimised model back to the original rotation files (or copies of them). optimised_rotation_updater.save_to_rotation_files() - + + # Generate post-run diagnostics: net rotation (median +/- MAD) and trench-normal + # migration (mean +/- MAD) through time, written as CSV and plots to 'model_output/'. + if generate_diagnostics: + from model_diagnostics import generate_model_diagnostics + generate_model_diagnostics( + datadir, + rotfile, + no_net_rotation_model.get_no_net_rotation_filename(), + trench_resolver, + age_range, + interval, + get_reference_params, + model_name) + main_end_time = round(time.time() - main_start, 10) main_sec = timedelta(seconds = float(main_end_time)) main_dt = datetime(1,1,1) + main_sec diff --git a/README.md b/README.md index 11eb74b..4dc2265 100644 --- a/README.md +++ b/README.md @@ -159,6 +159,19 @@ By default the workflow now *screens* all seeds before optimising. The objective An empirical convergence study on a 1Ga global plate model found that screening 400 seeds and fully optimising the best 16 (plus 16 uniform) reproduces the full 400-seed multistart minimum to within ~0.5% cost, while reducing the per-timestep optimisation cost by roughly a factor of 10. The study also found that the optimised pole is only loosely constrained near the minimum (several near-degenerate minima can differ by 10-20 degrees in pole position while differing by less than 0.5% in cost), so the residual error introduced by screening is well below the intrinsic uncertainty of the objective function itself. +Summary of the convergence study (two representative timesteps; full results in "model_output/seed_study/"): + +| Strategy | Full optimisations per timestep | Best cost vs 400-seed multistart | Approx. CPU cost | +|---|---|---|---| +| 400-seed multistart (original default) | 400 | - (reference) | 1.0x | +| 100-seed uniform grid | 100 | -0.2% / -0.3% | 0.25x | +| 49-seed uniform grid | 49 | +1.5% / -0.3% | 0.12x | +| screen all 400, polish best 16 + 16 uniform | ~30 | +0.5% / +0.3% | ~0.1x | + +(Negative values mean the reduced strategy found a slightly *better* minimum than the full multistart - the +multistart minimum itself is reproducible only to within ~0.5% because the optimiser converges to one of +several near-degenerate minima.) + To reproduce or extend the study (eg, for a different plate model or different seed counts), use the standalone, resumable harness "seed_study.py": ``` @@ -178,13 +191,22 @@ The original trench migration component drives trench-normal migration velocitie The default scheme is therefore now `trench_migration_scheme = 'rollback'` (in "Optimised_config.py"): per-trench orthogonal velocities are driven toward a target of +10 mm/yr retreat, and the bounds on the *mean* orthogonal velocity are tightened to (0, 20) mm/yr - the global mean must be retreating, but at less than 2 cm/yr. Individual trenches can still advance; only the mean is required to retreat. This makes trench kinematics an independent constraint that competes with the net rotation minimisation (instead of duplicating it). Set `trench_migration_scheme = 'minimise'` (or `OPTAPM_TM_SCHEME=minimise`) to restore the original behaviour. +A single-timestep validation (5 Ma) of the two schemes, showing that the rollback scheme makes trench kinematics an independent constraint (net rotation rises instead of being co-minimised, and the trench population becomes retreat-dominated as observed): + +| Scheme | Mean trench-normal velocity | Segments retreating | Net rotation at optimum | +|---|---|---|---| +| `minimise` (original) | -1.8 mm/yr (net advance) | 56% | 0.144 deg/Myr | +| `rollback` (new default) | +1.8 mm/yr (net retreat) | 63% | 0.175 deg/Myr | + +The parameter choices are: target retreat of +10 mm/yr = the centre of the observed median trench-normal retreat of +0.9 to +1.3 cm/yr across reference frames (Schellart et al. 2008); mean bounds of (0, 20) mm/yr = the global mean must retreat (62-78% of segments retreat in all reference frames examined by Schellart et al. 2008) but more slowly than ~2 cm/yr (just above the observed mean retreat of +1.3-1.5 cm/yr). + ### Uncertainty quantification (run variants) The optimisation is dominated by the net rotation (NR) minimisation: subduction zone migration largely depends on the net rotation optimisation (it is not an independent parameter), and the same holds for limiting the speed of continents (see [Muller et al. 2022, Solid Earth](https://doi.org/10.5194/se-13-1127-2022)). Perturbing the component weights only slightly changes the outcome, so a defensible uncertainty envelope for the optimised reference frame instead consists of end-members in the net rotation bounds: 1. **No-net-rotation (NNR) end-member** - zero net rotation. This rotation file is already produced by every run as "no_net_rotation_model_.rot". 2. **Best run** (`run_variant = 'best'`, the default) - NR bounded to (0.08, 0.20) deg/Myr: non-zero (as suggested by mantle flow models, eg, Becker 2006) but below the preferred geodynamic upper limit (Conrad & Behn 2010). -3. **Maximum-NR end-member** (`run_variant = 'nr_max'`) - the NR upper bound relaxed to 0.30 deg/Myr, the top of the range permitted by asthenospheric shear / seismic anisotropy constraints (Conrad & Behn 2010 derive NR < 0.26 deg/Myr for an asthenosphere 10x less viscous than the upper mantle; anisotropy proxies allow 0.2-0.3 deg/Myr). Larger values approach the Pacific hotspot (HS3) frame rates (0.33-0.44 deg/Myr), which are widely considered geodynamically implausible. +3. **Maximum-NR end-member** (`run_variant = 'nr_max'`) - the NR upper bound relaxed to 0.30 deg/Myr *and the NR weight halved* (inverse weight 2.0). The 0.30 deg/Myr ceiling is the top of the range permitted by asthenospheric shear / seismic anisotropy constraints (Conrad & Behn 2010 derive NR < 0.26 deg/Myr for an asthenosphere 10x less viscous than the upper mantle; anisotropy proxies allow 0.2-0.3 deg/Myr); larger values approach the Pacific hotspot (HS3) frame rates (0.33-0.44 deg/Myr), which are widely considered geodynamically implausible. The weight is halved because relaxing the bounds alone is not sufficient to produce a directional end-member: the bounds are hard penalty walls that only affect timesteps where the optimum presses against the 0.20 deg/Myr ceiling (mostly in deep time), whereas inside the walls the solution is set by the smooth trade-off between the NR cost and the other components. Halving the NR weight shifts that trade-off at every timestep. Single-timestep tests (5 Ma) confirm this: relaxing the bounds alone merely selected a different near-degenerate minimum (NR 0.123 deg/Myr), whereas bounds + halved weight produced the intended systematically stronger-rollback, higher-NR solution (NR 0.177 deg/Myr, trench retreat strengthening from +1.8 to +3.2 mm/yr, 68% of segments retreating). To produce the envelope, run the workflow twice (the `nr_max` variant appends its name to the model name, so the two runs write separate output files): @@ -195,6 +217,24 @@ To produce the envelope, run the workflow twice (the `nr_max` variant appends it The variant NR bounds are defined in `run_variant_nr_bounds` in "Optimised_config.py" (where further variants can be added). +### Model diagnostics (summary statistics and plots) + +At the end of every run (if `generate_diagnostics = True`, the default), the workflow computes per-timestep summary statistics of the optimised model and writes them to "model_output/": + +* `_diagnostics.csv` - per timestep: median and MAD (median absolute deviation) of the net rotation rate (deg/Myr, sampled at 1 Myr sub-steps within each timestep), and mean, median and MAD of the trench-normal migration velocities (mm/yr, positive = retreat) across all subduction zone segments, plus the percentage of segments retreating. +* `_net_rotation.png` - net rotation (median +/- MAD) through time, with the 0.20 deg/Myr preferred upper bound and the 0.26 deg/Myr Conrad & Behn (2010) limit marked. +* `_trench_migration.png` - trench-normal migration (mean +/- MAD) through time, with the observed present-day mean retreat range (+0.9 to +1.5 cm/yr, Schellart et al. 2008) marked. + +The same statistics and plots are also generated for the no-net-rotation model (outputs named `_NNR_*`). This shows how subduction zone migration behaves in an NNR world (the zero-NR end-member of the uncertainty envelope), and the NNR net rotation plot doubles as a sanity check (it should be ~zero at all times). + +The diagnostics can be (re)generated standalone on an existing optimised model without re-running the optimisation: + +``` + python model_diagnostics.py +``` + +(This uses the current settings in "Optimised_config.py", including any `OPTAPM_*` environment overrides, eg, `OPTAPM_VARIANT=nr_max python model_diagnostics.py` to diagnose the nr_max run.) + ### Quick test runs The age range, seed count, screening parameters and parallelisation can be temporarily overridden via environment variables without editing "Optimised_config.py" - useful for quick smoke tests: @@ -311,3 +351,14 @@ This can be done using the 'remove_plate_rotations' module in GPlately ( https:/ After running the workflow and post-processing (although you don't need to run `gplately.ptt.remove_plate_rotations` for this), you can plot the trench migration stats and net rotation curves for the non-optimised and any optimised models using the Jupyter notebooks in the `figures/` directory. + + +## References for parameter choices + +* Atkins, S., & Coltice, N. (2021). Constraining the range and variation of lithospheric net rotation using geodynamic modeling. *Journal of Geophysical Research: Solid Earth*, 126, e2021JB022057. https://doi.org/10.1029/2021JB022057 - convection models produce Earth-like net rotation mostly well below 0.2 deg/Myr. +* Becker, T. W. (2006). On the effect of temperature and strain-rate dependent viscosity on global mantle flow, net rotation, and plate-driving forces. *Geophysical Journal International*, 167, 943-957. https://doi.org/10.1111/j.1365-246X.2006.03172.x - mantle flow models suggest net rotation should be positive and non-zero (basis for the 0.08 deg/Myr lower bound). +* Conrad, C. P., & Behn, M. D. (2010). Constraints on lithosphere net rotation and asthenospheric viscosity from global mantle flow models and seismic anisotropy. *Geochemistry, Geophysics, Geosystems*, 11, Q05W05. https://doi.org/10.1029/2009GC002970 - net rotation < 0.26 deg/Myr for an asthenosphere 10x less viscous than the upper mantle (basis for the 0.20 deg/Myr preferred and 0.30 deg/Myr maximum upper bounds). +* Muller, R. D., et al. (2022). A tectonic-rules-based mantle reference frame since 1 billion years ago - implications for supercontinent cycles and plate-mantle system evolution. *Solid Earth*, 13, 1127-1159. https://doi.org/10.5194/se-13-1127-2022 - shows that net rotation minimisation dominates the optimisation and that trench migration is not an independent constraint under the original ('minimise') scheme (basis for the uncertainty-envelope design and the 'rollback' scheme). +* Schellart, W. P., Stegman, D. R., & Freeman, J. (2008). Global trench migration velocities and slab migration induced upper mantle volume fluxes: Constraints to find an Earth reference frame based on minimizing viscous dissipation. *Earth-Science Reviews*, 88, 118-144. https://doi.org/10.1016/j.earscirev.2008.01.005 - 62-78% of trench segments retreat across reference frames, with mean trench-normal retreat of +1.3-1.5 cm/yr and medians of +0.9-1.3 cm/yr (basis for the +10 mm/yr rollback target and the (0, 20) mm/yr mean bounds). +* Tetley, M. G., Williams, S. E., Gurnis, M., Flament, N., & Muller, R. D. (2019). Constraining absolute plate motions since the Triassic. *Journal of Geophysical Research: Solid Earth*, 124, 7231-7258. https://doi.org/10.1029/2019JB017442 - the original optAPM optimisation framework. +* Williams, S., Flament, N., Muller, R. D., & Butterworth, N. (2015). Absolute plate motions since 130 Ma constrained by subduction zone kinematics. *Earth and Planetary Science Letters*, 418, 66-77. https://doi.org/10.1016/j.epsl.2015.02.026 - trench kinematic plausibility (dominantly slow retreat, minimal advance) as a primary constraint on absolute plate motion. diff --git a/model_diagnostics.py b/model_diagnostics.py new file mode 100644 index 0000000..4adbdde --- /dev/null +++ b/model_diagnostics.py @@ -0,0 +1,278 @@ +""" +Post-run diagnostics for an optimised absolute plate motion model. + +Computes, for every timestep of an optimisation run: + + * Net rotation (NR) of the optimised model: the rate (deg/Myr) at 1 Myr sub-steps + within each timestep interval, summarised per timestep as the median and the + median absolute deviation (MAD) of those sub-step rates. + + * Trench-normal migration of all subduction zone segments: per timestep, the mean + and MAD of the trench-orthogonal velocities (mm/yr, positive = retreat/rollback) + across all trench segments, plus the percentage of segments retreating. + +Outputs (in 'model_output/'): + + * _diagnostics.csv - the statistics through time + * _net_rotation.png - NR median +/- MAD vs time + * _trench_migration.png - TM mean +/- MAD vs time + +These diagnostics run automatically at the end of 'Optimised_APM.py' (if +'generate_diagnostics' is True in 'Optimised_config.py'), and can also be run +standalone on an existing optimised model: + + python model_diagnostics.py + +(using the current settings in 'Optimised_config.py', including any OPTAPM_* +environment variable overrides such as OPTAPM_VARIANT / OPTAPM_START_AGE). +""" + +import math +import os.path +import sys + +import numpy as np +import pygplates + +import optimisation_methods +import subduction_convergence_for_absolute_plate_motion as scap + + +def mad(values, centre): + """Median absolute deviation about 'centre'.""" + values = np.asarray(values) + if values.size == 0: + return float('nan') + return float(np.median(np.abs(values - centre))) + + +def calculate_diagnostics( + data_dir, + optimised_rotation_filename, # Relative to the 'data/' directory. + no_net_rotation_filename, # Relative to the 'data/' directory. + trench_resolver, # TrenchResolver (used to resolve trenches at each time). + age_range, # Iterable of timestep start ages (eg, [5, 10, ..., 1000]). + interval, + reference_params_function): # Function (accepting age) returning (ref_plate_id, _). + """ + Calculate net rotation and trench migration statistics at each timestep. + Returns a list of dict rows (one per timestep age). + """ + + optimised_rotation_model = pygplates.RotationModel( + os.path.join(data_dir, optimised_rotation_filename)) + nn_rotation_model = pygplates.RotationModel( + os.path.join(data_dir, no_net_rotation_filename)) + + rows = [] + + for age in age_range: + + ref_rotation_plate_id, _ = reference_params_function(age) + + # + # Net rotation: rates at 1 Myr sub-steps within [age - interval, age]. + # + nr_timesteps = range(age - interval, age, 1) + _, _, PTangle1, _, _, _, _, _, _ = optimisation_methods.ApproximateNR_from_features( + optimised_rotation_model, nn_rotation_model, nr_timesteps, ref_rotation_plate_id) + nr_rates = np.abs(np.asarray(PTangle1, dtype=float)) # deg/Myr at each 1 Myr sub-step + nr_median = float(np.median(nr_rates)) + nr_mad = mad(nr_rates, nr_median) + + # + # Trench migration: trench-orthogonal velocities across all subduction segments. + # (Positive = retreat/rollback toward the ocean basin; negative = advance.) + # + trench_resolver.generate_resolved_trenches(age) + trench_features = pygplates.FeatureCollection( + os.path.join(data_dir, trench_resolver.get_trench_migration_filename())) + + tm_stats = scap.subduction_absolute_motion( + optimised_rotation_model, trench_features, np.radians(1.0), age, + velocity_delta_time=interval) + trench_vel = np.array([stat[2] for stat in tm_stats]) * 10 # cm/yr -> mm/yr + trench_obl = np.array([stat[3] for stat in tm_stats]) + tm_vel_orth = np.abs(trench_vel) * -np.cos(np.radians(trench_obl)) + + tm_mean = float(np.mean(tm_vel_orth)) + tm_median = float(np.median(tm_vel_orth)) + tm_mad = mad(tm_vel_orth, tm_median) + pct_retreating = float(100.0 * np.mean(tm_vel_orth > 0)) + + rows.append({ + 'age': age, + 'net_rotation_median_deg_per_myr': nr_median, + 'net_rotation_mad_deg_per_myr': nr_mad, + 'trench_migration_mean_mm_per_yr': tm_mean, + 'trench_migration_median_mm_per_yr': tm_median, + 'trench_migration_mad_mm_per_yr': tm_mad, + 'percent_trenches_retreating': pct_retreating, + }) + + print(' diagnostics at {0} Ma: NR median {1:.3f} deg/Myr, TM mean {2:+.1f} mm/yr ({3:.0f}% retreating)'.format( + age, nr_median, tm_mean, pct_retreating)) + sys.stdout.flush() + + return rows + + +def write_diagnostics_csv(rows, csv_filename): + fields = ['age', + 'net_rotation_median_deg_per_myr', 'net_rotation_mad_deg_per_myr', + 'trench_migration_mean_mm_per_yr', 'trench_migration_median_mm_per_yr', + 'trench_migration_mad_mm_per_yr', 'percent_trenches_retreating'] + with open(csv_filename, 'w') as f: + f.write(','.join(fields) + '\n') + for row in rows: + f.write(','.join('{0:.4f}'.format(row[field]) if field != 'age' else str(row[field]) + for field in fields) + '\n') + print('Wrote {0}'.format(csv_filename)) + + +def plot_diagnostics(rows, model_name, output_dir): + """Generate the net rotation and trench migration plots (PNG).""" + # Use a non-interactive backend so this works on HPC nodes without a display. + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + + ages = np.array([row['age'] for row in rows]) + + # + # Net rotation plot. + # + nr_median = np.array([row['net_rotation_median_deg_per_myr'] for row in rows]) + nr_mad = np.array([row['net_rotation_mad_deg_per_myr'] for row in rows]) + + fig, ax = plt.subplots(figsize=(10, 5)) + ax.fill_between(ages, nr_median - nr_mad, nr_median + nr_mad, + alpha=0.3, label='median +/- MAD (within timestep)') + ax.plot(ages, nr_median, lw=1.5, marker='o', ms=2.5, label='median') + # Reference lines: preferred geodynamic upper limit and zero. + ax.axhline(0.20, ls='--', lw=0.8, color='grey', + label='0.20 deg/Myr (preferred upper bound)') + ax.axhline(0.26, ls=':', lw=0.8, color='grey', + label='0.26 deg/Myr (Conrad & Behn 2010 limit)') + whole_median = float(np.median(nr_median)) + whole_mad = mad(nr_median, whole_median) + ax.set_xlabel('Age (Ma)') + ax.set_ylabel('Net rotation (deg/Myr)') + ax.set_title('{0}: net rotation through time\n' + '(whole-run median {1:.3f} +/- {2:.3f} deg/Myr MAD)'.format( + model_name, whole_median, whole_mad), fontsize=10) + ax.invert_xaxis() + ax.legend(fontsize=8) + ax.grid(alpha=0.3) + fig.tight_layout() + nr_png = os.path.join(output_dir, '{0}_net_rotation.png'.format(model_name)) + fig.savefig(nr_png, dpi=150) + plt.close(fig) + print('Wrote {0}'.format(nr_png)) + + # + # Trench migration plot. + # + tm_mean = np.array([row['trench_migration_mean_mm_per_yr'] for row in rows]) + tm_mad = np.array([row['trench_migration_mad_mm_per_yr'] for row in rows]) + pct_retreat = np.array([row['percent_trenches_retreating'] for row in rows]) + + fig, ax = plt.subplots(figsize=(10, 5)) + ax.fill_between(ages, tm_mean - tm_mad, tm_mean + tm_mad, + alpha=0.3, label='mean +/- MAD (across trench segments)') + ax.plot(ages, tm_mean, lw=1.5, marker='o', ms=2.5, label='mean') + ax.axhline(0.0, ls='-', lw=0.8, color='black') + # Observed present-day range of mean trench-normal retreat (Schellart et al. 2008). + ax.axhspan(9, 15, alpha=0.15, color='green', + label='observed present-day mean retreat\n(+0.9 to +1.5 cm/yr, Schellart et al. 2008)') + whole_mean = float(np.mean(tm_mean)) + ax.set_xlabel('Age (Ma)') + ax.set_ylabel('Trench-normal migration (mm/yr; positive = retreat)') + ax.set_title('{0}: subduction zone migration through time\n' + '(whole-run mean {1:+.1f} mm/yr; mean {2:.0f}% of segments retreating)'.format( + model_name, whole_mean, float(np.mean(pct_retreat))), fontsize=10) + ax.invert_xaxis() + ax.legend(fontsize=8) + ax.grid(alpha=0.3) + fig.tight_layout() + tm_png = os.path.join(output_dir, '{0}_trench_migration.png'.format(model_name)) + fig.savefig(tm_png, dpi=150) + plt.close(fig) + print('Wrote {0}'.format(tm_png)) + + +def generate_model_diagnostics( + data_dir, + optimised_rotation_filename, + no_net_rotation_filename, + trench_resolver, + age_range, + interval, + reference_params_function, + model_name, + output_dir='model_output', + also_diagnose_nnr_model=True): + """ + Compute the diagnostics, write the CSV and generate the two plots. + This is called at the end of a standard optimisation run (see 'Optimised_APM.py'). + + If 'also_diagnose_nnr_model' is True then the same diagnostics are also generated + for the no-net-rotation model itself (with outputs named '_NNR_*'). + This shows how subduction zone migration behaves in an NNR world - the zero-NR + end-member of the uncertainty envelope - and the NNR net rotation plot doubles + as a sanity check (it should be approximately zero at all times). + """ + print('Generating model diagnostics (net rotation and trench migration through time)...') + sys.stdout.flush() + + if not os.path.isdir(output_dir): + os.makedirs(output_dir) + + rows = calculate_diagnostics( + data_dir, optimised_rotation_filename, no_net_rotation_filename, + trench_resolver, age_range, interval, reference_params_function) + write_diagnostics_csv(rows, os.path.join(output_dir, '{0}_diagnostics.csv'.format(model_name))) + plot_diagnostics(rows, model_name, output_dir) + + if also_diagnose_nnr_model: + print('Generating diagnostics for the no-net-rotation model...') + sys.stdout.flush() + # Use the NNR model as the rotation model being diagnosed. Its net rotation + # (relative to itself) should be ~zero - a useful sanity check - while its + # trench migration statistics show subduction zone kinematics in an NNR world. + nnr_rows = calculate_diagnostics( + data_dir, no_net_rotation_filename, no_net_rotation_filename, + trench_resolver, age_range, interval, reference_params_function) + nnr_name = '{0}_NNR'.format(model_name) + write_diagnostics_csv(nnr_rows, os.path.join(output_dir, '{0}_diagnostics.csv'.format(nnr_name))) + plot_diagnostics(nnr_rows, nnr_name, output_dir) + + return rows + + +if __name__ == '__main__': + + # Standalone mode: run the diagnostics on an existing optimised model using + # the current settings in 'Optimised_config.py'. + from Optimised_config import (datadir, data_model, model_name, start_age, actual_end_age, + interval, get_reference_params, original_rotation_filenames, + topology_filenames) + from trench_resolver import TrenchResolver + + optimised_rotation_filename = os.path.join( + data_model, 'optimisation', 'optimised_rotation_model_' + model_name + '.rot') + no_net_rotation_filename = os.path.join( + data_model, 'optimisation', 'no_net_rotation_model_' + model_name + '.rot') + + print('Loading topologies...') + topology_features = [] + for topology_filename in topology_filenames: + topology_features.extend(pygplates.FeatureCollection(os.path.join(datadir, topology_filename))) + + trench_resolver = TrenchResolver(datadir, original_rotation_filenames, topology_features, data_model) + + age_range = range(actual_end_age + interval, start_age + interval, interval) + + generate_model_diagnostics( + datadir, optimised_rotation_filename, no_net_rotation_filename, + trench_resolver, age_range, interval, get_reference_params, model_name)