Repository: plasmacontrol/desc
Issue: Constraints have wrong derivatives when building at initialization

Any `Objective` can be passed an `Equilibrium` when initialized to build the objective, or otherwise it can be built later.  The former method gives incorrect derivative values.  This has only been tested with linear objectives that are typically used for constraints (specifically, the problem seems to only be with `FixBoundaryR` and `FixBoundaryZ`).  

Minimal working example to reproduce the bug:
```
import numpy as np
from desc.equilibrium import Equilibrium
from desc.objectives import FixBoundaryR, FixBoundaryZ

eq = Equilibrium(M=3, N=1)

# initialize the constraints without building
constraints_1 = (
    FixBoundaryR(fixed_boundary=False),
    FixBoundaryZ(fixed_boundary=False),
    FixBoundaryR(fixed_boundary=True),
    FixBoundaryZ(fixed_boundary=True),
)
print("method 1: without building\\n")
for obj in constraints_1:
    obj.build(eq)
    arg = obj.args[0]
    A = obj.derivatives["jac"][arg](np.zeros(obj.dimensions[arg]))
    print("constraint {}: A.shape = {}, max(|A|) = {}".format(
        obj.name, A.shape, np.max(np.abs(A))))

# build the constraints while initializing
constraints_2 = (
    FixBoundaryR(fixed_boundary=False, eq=eq),
    FixBoundaryZ(fixed_boundary=False, eq=eq),
    FixBoundaryR(fixed_boundary=True, eq=eq),
    FixBoundaryZ(fixed_boundary=True, eq=eq),
)
print("\\nmethod 2: with building\\n")
for obj in constraints_2:
    arg = obj.args[0]
    A = obj.derivatives["jac"][arg](np.zeros(obj.dimensions[arg]))
    print("constraint {}: A.shape = {}, max(|A|) = {}".format(
        obj.name, A.shape, np.max(np.abs(A))))
```
And the output: 
```
method 1: without building

constraint lcfs R: A.shape = (21, 21), max(|A|) = 1.0
constraint lcfs Z: A.shape = (21, 21), max(|A|) = 1.0
constraint lcfs R: A.shape = (21, 30), max(|A|) = 1.0
constraint lcfs Z: A.shape = (21, 30), max(|A|) = 1.0

method 2: with building

constraint lcfs R: A.shape = (21, 21), max(|A|) = 0.0
constraint lcfs Z: A.shape = (21, 21), max(|A|) = 0.0
constraint lcfs R: A.shape = (21, 21), max(|A|) = 0.0
constraint lcfs Z: A.shape = (21, 21), max(|A|) = 0.0
```
"method 1" is correct.  "method 2" gives derivative matrices that are all zeros, and they are the wrong shape for `fixed_boundary=True`.  


Changed files:
desc/objectives/objective_funs.py, tests/test_linear_objectives.py, desc/objectives/linear_objectives.py

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index ba10c021..7dbb6e25 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -75,6 +75,7 @@ class FixBoundaryR(_Objective):
         self._fixed_boundary = fixed_boundary
         self._modes = modes
         self._surface_label = surface_label
+        self._args = ["R_lmn"] if self._fixed_boundary else ["Rb_lmn"]
         super().__init__(
             eq=eq,
             target=target,
@@ -83,8 +84,6 @@ class FixBoundaryR(_Objective):
             normalize_target=normalize_target,
             name=name,
         )
-        self._print_value_fmt = "R boundary error: {:10.3e} (m)"
-        self._args = ["R_lmn"] if self._fixed_boundary else ["Rb_lmn"]
 
     def build(self, eq, use_jit=True, verbose=1):
         """Build constant arrays.
@@ -224,6 +223,7 @@ class FixBoundaryZ(_Objective):
         self._fixed_boundary = fixed_boundary
         self._modes = modes
         self._surface_label = surface_label
+        self._args = ["Z_lmn"] if self._fixed_boundary else ["Zb_lmn"]
         super().__init__(
             eq=eq,
             target=target,
@@ -232,8 +232,6 @@ class FixBoundaryZ(_Objective):
             normalize_target=normalize_target,
             name=name,
         )
-        self._print_value_fmt = "Z boundary error: {:10.3e} (m)"
-        self._args = ["Z_lmn"] if self._fixed_boundary else ["Zb_lmn"]
 
     def build(self, eq, use_jit=True, verbose=1):
         """Build constant arrays.
diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py
index 6d32b1cb..2e789f47 100644
--- a/desc/objectives/objective_funs.py
+++ b/desc/objectives/objective_funs.py
@@ -489,7 +489,12 @@ class _Objective(IOAble, ABC):
         self._name = name
         self._use_jit = None
         self._built = False
-        self._args = [arg for arg in getfullargspec(self.compute)[0] if arg != "self"]
+        # if args is already set don't overwrite it
+        self._args = getattr(
+            self,
+            "_args",
+            [arg for arg in getfullargspec(self.compute)[0] if arg != "self"],
+        )
         if eq is not None:
             self.build(eq)
 
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index 644af70b..a9a916f3 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -5,7 +5,13 @@ import pytest
 
 from desc.equilibrium import Equilibrium
 from desc.geometry import FourierRZToroidalSurface
-from desc.objectives import FixBoundaryZ, FixCurrent, FixIota, FixLambdaGauge
+from desc.objectives import (
+    FixBoundaryR,
+    FixBoundaryZ,
+    FixCurrent,
+    FixIota,
+    FixLambdaGauge,
+)
 from desc.profiles import PowerSeriesProfile
 
 
@@ -112,3 +118,65 @@ def test_constrain_asserts():
     eq = Equilibrium(current=PowerSeriesProfile(0))
     with pytest.raises(ValueError):
         eq.solve(constraints=(FixCurrent(), FixIota()))
+
+
+@pytest.mark.unit
+def test_build_init():
+    """Ensure that passing an equilibrium to init builds the objective correctly.
+
+    Related to gh issue #378
+    """
+    eq = Equilibrium(M=3, N=1)
+
+    # initialize the constraints without building
+    fbR1 = FixBoundaryR(fixed_boundary=False)
+    fbZ1 = FixBoundaryZ(fixed_boundary=False)
+    fbR2 = FixBoundaryR(fixed_boundary=True)
+    fbZ2 = FixBoundaryZ(fixed_boundary=True)
+    for obj in (fbR1, fbR2, fbZ1, fbZ2):
+        obj.build(eq)
+
+    arg = fbR1.args[0]
+    A = fbR1.derivatives["jac"][arg](np.zeros(fbR1.dimensions[arg]))
+    assert np.max(np.abs(A)) == 1
+    assert A.shape == (eq.surface.R_basis.num_modes, eq.surface.R_basis.num_modes)
+
+    arg = fbR2.args[0]
+    A = fbR2.derivatives["jac"][arg](np.zeros(fbR2.dimensions[arg]))
+    assert np.max(np.abs(A)) == 1
+    assert A.shape == (eq.surface.R_basis.num_modes, eq.R_basis.num_modes)
+
+    arg = fbZ1.args[0]
+    A = fbZ1.derivatives["jac"][arg](np.zeros(fbZ1.dimensions[arg]))
+    assert np.max(np.abs(A)) == 1
+    assert A.shape == (eq.surface.Z_basis.num_modes, eq.surface.Z_basis.num_modes)
+
+    arg = fbZ2.args[0]
+    A = fbZ2.derivatives["jac"][arg](np.zeros(fbZ2.dimensions[arg]))
+    assert np.max(np.abs(A)) == 1
+    assert A.shape == (eq.surface.Z_basis.num_modes, eq.Z_basis.num_modes)
+
+    fbR1 = FixBoundaryR(fixed_boundary=False, eq=eq)
+    fbZ1 = FixBoundaryZ(fixed_boundary=False, eq=eq)
+    fbR2 = FixBoundaryR(fixed_boundary=True, eq=eq)
+    fbZ2 = FixBoundaryZ(fixed_boundary=True, eq=eq)
+
+    arg = fbR1.args[0]
+    A = fbR1.derivatives["jac"][arg](np.zeros(fbR1.dimensions[arg]))
+    assert np.max(np.abs(A)) == 1
+    assert A.shape == (eq.surface.R_basis.num_modes, eq.surface.R_basis.num_modes)
+
+    arg = fbR2.args[0]
+    A = fbR2.derivatives["jac"][arg](np.zeros(fbR2.dimensions[arg]))
+    assert np.max(np.abs(A)) == 1
+    assert A.shape == (eq.surface.R_basis.num_modes, eq.R_basis.num_modes)
+
+    arg = fbZ1.args[0]
+    A = fbZ1.derivatives["jac"][arg](np.zeros(fbZ1.dimensions[arg]))
+    assert np.max(np.abs(A)) == 1
+    assert A.shape == (eq.surface.Z_basis.num_modes, eq.surface.Z_basis.num_modes)
+
+    arg = fbZ2.args[0]
+    A = fbZ2.derivatives["jac"][arg](np.zeros(fbZ2.dimensions[arg]))
+    assert np.max(np.abs(A)) == 1
+    assert A.shape == (eq.surface.Z_basis.num_modes, eq.Z_basis.num_modes)