From 1e3141b9ef093eec2b0261f1a674aca0a85027c4 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Mon, 30 Sep 2024 16:37:11 -0400 Subject: [PATCH 1/5] Update rules for chunking jacobian --- desc/objectives/objective_funs.py | 48 +++++++++++++++---------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 9cdfb61c24..88ae520e4b 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -113,14 +113,6 @@ def __init__( self._compiled = False self._name = name - def _set_derivatives(self): - """Choose derivative mode based on mode of sub-objectives.""" - if self._deriv_mode == "auto": - if all((obj._deriv_mode == "fwd") for obj in self.objectives): - self._deriv_mode = "batched" - else: - self._deriv_mode = "blocked" - def _unjit(self): """Remove jit compiled methods.""" methods = [ @@ -178,12 +170,20 @@ def build(self, use_jit=None, verbose=1): else: self._scalar = False - self._set_derivatives() + self._set_things() + + # setting derivative mode and chunking. + errorif( + isposint(self._jac_chunk_size) and self._deriv_mode in ["auto", "blocked"], + ValueError, + "'jac_chunk_size' was passed into ObjectiveFunction, but the " + "ObjectiveFunction is not using 'batched' deriv_mode", + ) sub_obj_jac_chunk_sizes_are_ints = [ isposint(obj._jac_chunk_size) for obj in self.objectives ] errorif( - any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode != "blocked", + any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode == "batched", ValueError, "'jac_chunk_size' was passed into one or more sub-objectives, but the" " ObjectiveFunction is using 'batched' deriv_mode, so sub-objective " @@ -192,21 +192,12 @@ def build(self, use_jit=None, verbose=1): " Specify 'blocked' deriv_mode if each sub-objective is desired to have a " "different 'jac_chunk_size' for its Jacobian computation.", ) - errorif( - self._jac_chunk_size not in ["auto", None] - and self._deriv_mode == "blocked", - ValueError, - "'jac_chunk_size' was passed into ObjectiveFunction, but the " - "ObjectiveFunction is using 'blocked' deriv_mode, so sub-objective " - "'jac_chunk_size' are used to compute each sub-objective's Jacobian, " - "`ignoring the ObjectiveFunction's 'jac_chunk_size'.", - ) - - if not self.use_jit: - self._unjit() - self._set_things() - self._built = True + if self._deriv_mode == "auto": + if all((obj._deriv_mode == "fwd") for obj in self.objectives): + self._deriv_mode = "batched" + else: + self._deriv_mode = "blocked" if self._jac_chunk_size == "auto": # Heuristic estimates of fwd mode Jacobian memory usage, @@ -218,6 +209,15 @@ def build(self, use_jit=None, verbose=1): * self.dim_x ) self._jac_chunk_size = max([1, max_chunk_size]) + if self._deriv_mode == "blocked": + for obj in self.objectives: + if obj._jac_chunk_size is None: + obj._jac_chunk_size = self._jac_chunk_size + + if not self.use_jit: + self._unjit() + + self._built = True timer.stop("Objective build") if verbose > 1: From 04121b18905630a1b91bb9be206b1d04fd2793f9 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Tue, 1 Oct 2024 15:27:11 -0400 Subject: [PATCH 2/5] Allow accessing dim_f during objective build --- desc/objectives/objective_funs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 88ae520e4b..ae1e294206 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -799,7 +799,7 @@ def dim_x(self): @property def dim_f(self): """int: Number of objective equations.""" - if not self.built: + if not hasattr(self, "_dim_f"): raise RuntimeError("ObjectiveFunction must be built first.") return self._dim_f From 3e3a07086b61db366afc0ba184589e376fc21839 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Tue, 1 Oct 2024 16:02:44 -0400 Subject: [PATCH 3/5] Use compatible deriv mode if jac_chunk_size is specified in getters --- desc/objectives/getters.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py index 727c6f1835..7d4772a8b0 100644 --- a/desc/objectives/getters.py +++ b/desc/objectives/getters.py @@ -1,6 +1,6 @@ """Utilities for getting standard groups of objectives and constraints.""" -from desc.utils import flatten_list, is_any_instance, unique_list +from desc.utils import flatten_list, is_any_instance, isposint, unique_list from ._equilibrium import Energy, ForceBalance, HelicalForceBalance, RadialForceBalance from .linear_objectives import ( @@ -86,7 +86,10 @@ def get_equilibrium_objective(eq, mode="force", normalize=True, jac_chunk_size=" objectives = (RadialForceBalance(**kwargs), HelicalForceBalance(**kwargs)) else: raise ValueError("got an unknown equilibrium objective type '{}'".format(mode)) - return ObjectiveFunction(objectives, jac_chunk_size=jac_chunk_size) + deriv_mode = "batched" if isposint(jac_chunk_size) else "auto" + return ObjectiveFunction( + objectives, jac_chunk_size=jac_chunk_size, deriv_mode=deriv_mode + ) def get_fixed_axis_constraints(eq, profiles=True, normalize=True): From 195753d8e79a30cce748bc645fd640b4eddbfe89 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Tue, 1 Oct 2024 20:52:16 -0400 Subject: [PATCH 4/5] Update tests for new chunking rules --- tests/test_examples.py | 1 + tests/test_objective_funs.py | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 53fb6479b5..93c700b36c 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1043,6 +1043,7 @@ def test_freeb_vacuum(): objective = ObjectiveFunction( VacuumBoundaryError(eq=eq, field=ext_field, field_fixed=True), jac_chunk_size=1000, + deriv_mode="batched", ) eq, _ = eq.optimize( objective, diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index cfbf09090b..0413a64926 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -1323,6 +1323,7 @@ def test_derivative_modes(): AspectRatio(eq), ], deriv_mode="batched", + jac_chunk_size="auto", use_jit=False, ) obj2 = ObjectiveFunction( @@ -1332,13 +1333,15 @@ def test_derivative_modes(): AspectRatio(eq, jac_chunk_size=None), ], deriv_mode="blocked", + jac_chunk_size="auto", use_jit=False, ) obj1.build() obj2.build() # check that default size works for blocked - assert obj2.objectives[1]._jac_chunk_size is None - assert obj2.objectives[2]._jac_chunk_size is None + assert obj2.objectives[0]._jac_chunk_size == 2 + assert obj2.objectives[1]._jac_chunk_size > 0 + assert obj2.objectives[2]._jac_chunk_size > 0 # hard to say what size auto will give, just check it is >0 assert obj1._jac_chunk_size > 0 obj3.build() From b7a0558d5a8514a2ac9db008961159e47c925ed5 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 2 Oct 2024 22:40:55 -0400 Subject: [PATCH 5/5] Update desc/objectives/objective_funs.py Co-authored-by: Yigit Gunsur Elmacioglu <102380275+YigitElma@users.noreply.github.com> --- desc/objectives/objective_funs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index ae1e294206..e66bfdbde5 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -186,7 +186,7 @@ def build(self, use_jit=None, verbose=1): any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode == "batched", ValueError, "'jac_chunk_size' was passed into one or more sub-objectives, but the" - " ObjectiveFunction is using 'batched' deriv_mode, so sub-objective " + " ObjectiveFunction is using 'batched' deriv_mode, so sub-objective " "'jac_chunk_size' will be ignored in favor of the ObjectiveFunction's " f"'jac_chunk_size' of {self._jac_chunk_size}." " Specify 'blocked' deriv_mode if each sub-objective is desired to have a "