Merge branch 'master' into add_all_limits

desc/io/hdf5_io.py

@@ -349,17 +349,24 @@ def write_dict(self, thedict, where=None):
         loc = self.resolve_where(where)
         loc.create_dataset("__class__", data="dict")
         for key in thedict.keys():
-            try:
-                data = thedict[key]
-                compression = (
-                    "gzip"
-                    if isinstance(data, np.ndarray) and np.asarray(data).size > 1
-                    else None
-                )
-                loc.create_dataset(key, data=data, compression=compression)
-            except TypeError:
+            if isinstance(thedict[key], list):
                 group = loc.create_group(key)
-                self.write_obj(thedict[key], group)
+                self.write_list(thedict[key], where=group)
+            elif isinstance(thedict[key], dict):
+                group = loc.create_group(key)
+                self.write_dict(thedict[key], where=group)
+            else:
+                try:
+                    data = thedict[key]
+                    compression = (
+                        "gzip"
+                        if isinstance(data, np.ndarray) and np.asarray(data).size > 1
+                        else None
+                    )
+                    loc.create_dataset(key, data=data, compression=compression)
+                except TypeError:
+                    group = loc.create_group(key)
+                    self.write_obj(thedict[key], group)
 
     def write_list(self, thelist, where=None):
         """Write list to file in group specified by where argument.
@@ -375,14 +382,21 @@ def write_list(self, thelist, where=None):
         loc = self.resolve_where(where)
         loc.create_dataset("__class__", data="list")
         for i in range(len(thelist)):
-            try:
-                data = thelist[i]
-                compression = (
-                    "gzip"
-                    if isinstance(data, np.ndarray) and np.asarray(data).size > 1
-                    else None
-                )
-                loc.create_dataset(str(i), data=data, compression=compression)
-            except TypeError:
+            if isinstance(thelist[i], list):
                 subloc = loc.create_group(str(i))
-                self.write_obj(thelist[i], where=subloc)
+                self.write_list(thelist[i], where=subloc)
+            elif isinstance(thelist[i], dict):
+                subloc = loc.create_group(str(i))
+                self.write_dict(thelist[i], where=subloc)
+            else:
+                try:
+                    data = thelist[i]
+                    compression = (
+                        "gzip"
+                        if isinstance(data, np.ndarray) and np.asarray(data).size > 1
+                        else None
+                    )
+                    loc.create_dataset(str(i), data=data, compression=compression)
+                except TypeError:
+                    subloc = loc.create_group(str(i))
+                    self.write_obj(thelist[i], where=subloc)

tests/benchmarks/benchmark_cpu_large.py

@@ -251,11 +251,12 @@ def test_perturb_1(SOLOVEV, benchmark):
 
     def setup():
         eq = EquilibriaFamily.load(load_from=str(SOLOVEV["desc_h5_path"]))[-1]
-        objective = get_equilibrium_objective()
-        constraints = get_fixed_boundary_constraints()
+        objective = get_equilibrium_objective(eq)
+        objective.build(eq)
+
+        constraints = get_fixed_boundary_constraints(eq)
         tr_ratio = [0.01, 0.25, 0.25]
         dp = np.zeros_like(eq.p_l)
-        objective.build(eq)
         dp[np.array([0, 2])] = 8e3 * np.array([1, -1])
         deltas = {"p_l": dp}
         args = (
@@ -283,11 +284,12 @@ def test_perturb_2(SOLOVEV, benchmark):
 
     def setup():
         eq = EquilibriaFamily.load(load_from=str(SOLOVEV["desc_h5_path"]))[-1]
-        objective = get_equilibrium_objective()
-        constraints = get_fixed_boundary_constraints()
+        objective = get_equilibrium_objective(eq)
+        objective.build(eq)
+
+        constraints = get_fixed_boundary_constraints(eq)
         tr_ratio = [0.01, 0.25, 0.25]
         dp = np.zeros_like(eq.p_l)
-        objective.build(eq)
         dp[np.array([0, 2])] = 8e3 * np.array([1, -1])
         deltas = {"p_l": dp}
         args = (

tests/benchmarks/benchmark_cpu_small.py

@@ -118,7 +118,7 @@ def test_objective_compile_heliotron(benchmark):
 
     def setup():
         eq = desc.examples.get("HELIOTRON")
-        objective = get_equilibrium_objective()
+        objective = get_equilibrium_objective(eq)
         objective.build(eq)
         args = (
             objective,
@@ -141,7 +141,7 @@ def test_objective_compile_dshape_current(benchmark):
 
     def setup():
         eq = desc.examples.get("DSHAPE_current")
-        objective = get_equilibrium_objective()
+        objective = get_equilibrium_objective(eq)
         objective.build(eq)
         args = (
             objective,
@@ -164,12 +164,9 @@ def test_objective_compile_atf(benchmark):
 
     def setup():
         eq = desc.examples.get("ATF")
-        objective = get_equilibrium_objective()
+        objective = get_equilibrium_objective(eq)
         objective.build(eq)
-        args = (
-            objective,
-            eq,
-        )
+        args = (objective, eq)
         kwargs = {}
         return args, kwargs
 
@@ -185,7 +182,7 @@ def run(objective, eq):
 def test_objective_compute_heliotron(benchmark):
     """Benchmark computing objective."""
     eq = desc.examples.get("HELIOTRON")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -202,7 +199,7 @@ def run(x):
 def test_objective_compute_dshape_current(benchmark):
     """Benchmark computing objective."""
     eq = desc.examples.get("DSHAPE_current")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -219,7 +216,7 @@ def run(x):
 def test_objective_compute_atf(benchmark):
     """Benchmark computing objective."""
     eq = desc.examples.get("ATF")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -236,7 +233,7 @@ def run(x):
 def test_objective_jac_heliotron(benchmark):
     """Benchmark computing jacobian."""
     eq = desc.examples.get("HELIOTRON")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -253,7 +250,7 @@ def run(x):
 def test_objective_jac_dshape_current(benchmark):
     """Benchmark computing jacobian."""
     eq = desc.examples.get("DSHAPE_current")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -270,7 +267,7 @@ def run(x):
 def test_objective_jac_atf(benchmark):
     """Benchmark computing jacobian."""
     eq = desc.examples.get("ATF")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -289,11 +286,11 @@ def test_perturb_1(benchmark):
 
     def setup():
         eq = desc.examples.get("SOLOVEV")
-        objective = get_equilibrium_objective()
-        constraints = get_fixed_boundary_constraints()
+        objective = get_equilibrium_objective(eq)
+        objective.build(eq)
+        constraints = get_fixed_boundary_constraints(eq)
         tr_ratio = [0.01, 0.25, 0.25]
         dp = np.zeros_like(eq.p_l)
-        objective.build(eq)
         dp[np.array([0, 2])] = 8e3 * np.array([1, -1])
         deltas = {"p_l": dp}
 
@@ -322,11 +319,11 @@ def test_perturb_2(benchmark):
 
     def setup():
         eq = desc.examples.get("SOLOVEV")
-        objective = get_equilibrium_objective()
-        constraints = get_fixed_boundary_constraints()
+        objective = get_equilibrium_objective(eq)
+        objective.build(eq)
+        constraints = get_fixed_boundary_constraints(eq)
         tr_ratio = [0.01, 0.25, 0.25]
         dp = np.zeros_like(eq.p_l)
-        objective.build(eq)
         dp[np.array([0, 2])] = 8e3 * np.array([1, -1])
         deltas = {"p_l": dp}
 

tests/benchmarks/benchmark_gpu_large.py

@@ -251,11 +251,12 @@ def test_perturb_1(SOLOVEV, benchmark):
 
     def setup():
         eq = EquilibriaFamily.load(load_from=str(SOLOVEV["desc_h5_path"]))[-1]
-        objective = get_equilibrium_objective()
-        constraints = get_fixed_boundary_constraints()
+        objective = get_equilibrium_objective(eq)
+        objective.build(eq)
+
+        constraints = get_fixed_boundary_constraints(eq)
         tr_ratio = [0.01, 0.25, 0.25]
         dp = np.zeros_like(eq.p_l)
-        objective.build(eq)
         dp[np.array([0, 2])] = 8e3 * np.array([1, -1])
         deltas = {"p_l": dp}
 
@@ -284,11 +285,12 @@ def test_perturb_2(SOLOVEV, benchmark):
 
     def setup():
         eq = EquilibriaFamily.load(load_from=str(SOLOVEV["desc_h5_path"]))[-1]
-        objective = get_equilibrium_objective()
-        constraints = get_fixed_boundary_constraints()
+        objective = get_equilibrium_objective(eq)
+        objective.build(eq)
+
+        constraints = get_fixed_boundary_constraints(eq)
         tr_ratio = [0.01, 0.25, 0.25]
         dp = np.zeros_like(eq.p_l)
-        objective.build(eq)
         dp[np.array([0, 2])] = 8e3 * np.array([1, -1])
         deltas = {"p_l": dp}
 

tests/benchmarks/benchmark_gpu_small.py

@@ -118,7 +118,7 @@ def test_objective_compile_heliotron(benchmark):
 
     def setup():
         eq = desc.examples.get("HELIOTRON")
-        objective = get_equilibrium_objective()
+        objective = get_equilibrium_objective(eq)
         objective.build(eq)
         args = (
             objective,
@@ -141,7 +141,7 @@ def test_objective_compile_dshape_current(benchmark):
 
     def setup():
         eq = desc.examples.get("DSHAPE_current")
-        objective = get_equilibrium_objective()
+        objective = get_equilibrium_objective(eq)
         objective.build(eq)
         args = (
             objective,
@@ -164,7 +164,7 @@ def test_objective_compile_atf(benchmark):
 
     def setup():
         eq = desc.examples.get("ATF")
-        objective = get_equilibrium_objective()
+        objective = get_equilibrium_objective(eq)
         objective.build(eq)
         args = (
             objective,
@@ -185,7 +185,7 @@ def run(objective, eq):
 def test_objective_compute_heliotron(benchmark):
     """Benchmark computing objective."""
     eq = desc.examples.get("HELIOTRON")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -202,7 +202,7 @@ def run(x):
 def test_objective_compute_dshape_current(benchmark):
     """Benchmark computing objective."""
     eq = desc.examples.get("DSHAPE_current")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -219,7 +219,7 @@ def run(x):
 def test_objective_compute_atf(benchmark):
     """Benchmark computing objective."""
     eq = desc.examples.get("ATF")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -236,7 +236,7 @@ def run(x):
 def test_objective_jac_heliotron(benchmark):
     """Benchmark computing jacobian."""
     eq = desc.examples.get("HELIOTRON")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -253,7 +253,7 @@ def run(x):
 def test_objective_jac_dshape_current(benchmark):
     """Benchmark computing jacobian."""
     eq = desc.examples.get("DSHAPE_current")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -270,7 +270,7 @@ def run(x):
 def test_objective_jac_atf(benchmark):
     """Benchmark computing jacobian."""
     eq = desc.examples.get("ATF")
-    objective = get_equilibrium_objective()
+    objective = get_equilibrium_objective(eq)
     objective.build(eq)
     objective.compile()
     x = objective.x(eq)
@@ -289,11 +289,11 @@ def test_perturb_1(benchmark):
 
     def setup():
         eq = desc.examples.get("SOLOVEV")
-        objective = get_equilibrium_objective()
-        constraints = get_fixed_boundary_constraints()
+        objective = get_equilibrium_objective(eq)
+        objective.build(eq)
+        constraints = get_fixed_boundary_constraints(eq)
         tr_ratio = [0.01, 0.25, 0.25]
         dp = np.zeros_like(eq.p_l)
-        objective.build(eq)
         dp[np.array([0, 2])] = 8e3 * np.array([1, -1])
         deltas = {"p_l": dp}
 
@@ -322,11 +322,11 @@ def test_perturb_2(benchmark):
 
     def setup():
         eq = desc.examples.get("SOLOVEV")
-        objective = get_equilibrium_objective()
-        constraints = get_fixed_boundary_constraints()
+        objective = get_equilibrium_objective(eq)
+        objective.build(eq)
+        constraints = get_fixed_boundary_constraints(eq)
         tr_ratio = [0.01, 0.25, 0.25]
         dp = np.zeros_like(eq.p_l)
-        objective.build(eq)
         dp[np.array([0, 2])] = 8e3 * np.array([1, -1])
         deltas = {"p_l": dp}