Modal idf (#360)

* fix fun3d imports

* update drivers

* updated trim sizing code

* update the oneway struct driver with correct derivatives now

* update multi driver

* fix zero deriv issues

* fix multi driver mistake

* remove second zero derivatives thing

* update trim sizing

* add set design parameter option

* Initial modal IDF commits (#359)

* Create modal_idf_driver.py

* Initial modal IDF commit

* adding modal analysis code from TACS

---------

Co-authored-by: Brian Burke <[email protected]>

funtofem/driver/__init__.py

@@ -12,6 +12,15 @@
 from .funtofem_nlbgs_driver import *
 from .funtofem_nlbgs_fsi_subiters_driver import *
 from .transfer_settings import *
-from .funtofem_shape_driver import *
+import importlib
+caps_loader = importlib.util.find_spec("pyCAPS")
+if caps_loader is not None:
+    from .funtofem_shape_driver import *
 from .oneway_struct_driver import *
 from .oneway_aero_driver import *
+
+# modal IDF driver
+from .modal_idf_driver import *
+
+# import all the custom or special drivers
+from .custom import *

funtofem/driver/_funtofem_driver.py

@@ -70,6 +70,14 @@ def __init__(
             whether to save and reload funtofem states
         """
 
+        # assert at least one coupled scenario
+        self.coupled_scenarios = []
+        if model is not None: # only case where model is None is fakeModel?
+            any_coupled = any([scenario.coupled for scenario in model.scenarios])
+            assert any_coupled
+
+            self.coupled_scenarios = [scenario for scenario in model.scenarios if scenario.coupled]
+
         # add the comm manger
         if comm_manager is not None:
             comm_manager = comm_manager
@@ -167,7 +175,8 @@ def solve_forward(self, steps=None):
             body.update_shape(complex_run)
 
         # loop over the forward problem for the different scenarios
-        for scenario in self.model.scenarios:
+        for scenario in self.coupled_scenarios:
+
             # tell the solvers what the variable values and functions are for this scenario
             if not self.fakemodel:
                 self._distribute_variables(scenario, self.model.bodies)
@@ -223,11 +232,10 @@ def solve_adjoint(self):
 
         # Zero the derivative values stored in the function
         self._zero_derivatives()
-        for func in functions:
-            func.zero_derivatives()
 
         # Set the functions into the solvers
-        for scenario in self.model.scenarios:
+        for scenario in self.coupled_scenarios:
+
             # tell the solvers what the variable values and functions are for this scenario
             self._distribute_variables(scenario, self.model.bodies)
             self._distribute_functions(scenario, self.model.bodies)
@@ -288,9 +296,11 @@ def _initialize_adjoint(self, scenario, bodies):
 
     def _zero_derivatives(self):
         """zero all model derivatives"""
-        for func in self.model.get_functions(all=True):
-            for var in self.model.get_variables():
-                func.derivatives[var] = 0.0
+        # TODO : only zero derivatives in coupled scenarios when using
+        for scenario in self.coupled_scenarios: # no need to zero composite functions, they are exactly diff later with no += effects
+            for func in scenario.functions:
+                for var in self.model.get_variables():
+                    func.derivatives[var] = 0.0
         return
 
     def _post_forward(self, scenario, bodies):

funtofem/driver/custom/__init__.py

@@ -0,0 +1,4 @@
+import importlib
+
+from .multi_driver import *
+from .oneway_struct_trim_driver import *

funtofem/driver/custom/multi_driver.py

@@ -0,0 +1,42 @@
+
+__all__ = ["MultiDriver"]
+
+class MultiDriver:
+    def __init__(self, driver_list:list):
+        """
+        call solve_forward, solve_adjoint on multiple drivers
+
+        useful if one scenario uses a coupled driver and the
+        other uses an uncoupled driver
+
+        in this way, we can include a mixture of each by combining these
+        drivers into one and still using the OptimizationManager
+        """
+        self.driver_list = driver_list
+
+        # copy comm and solvers from the first driver
+        # since these are used in optimizationManager
+        first_driver = self.driver_list[0]
+        self.comm = first_driver.comm
+        self.solvers = first_driver.solvers
+        self.model = first_driver.model
+
+    def solve_forward(self):
+        driver_list = self.driver_list
+        for driver in driver_list:
+            driver.solve_forward()
+
+    def solve_adjoint(self):
+        self._zero_derivatives()
+        driver_list = self.driver_list
+        for driver in driver_list:
+            driver.solve_adjoint()
+
+    def _zero_derivatives(self):
+        """zero all model derivatives"""
+        # TODO : only zero derivatives in coupled scenarios when using
+        model = self.driver_list[0].model
+        for func in model.get_functions(all=True):
+            for var in self.model.get_variables():
+                func.derivatives[var] = 0.0
+        return

funtofem/driver/custom/oneway_struct_trim_driver.py

@@ -0,0 +1,209 @@
+
+__all__ = ["OnewayStructTrimDriver"]
+
+from ..oneway_struct_driver import OnewayStructDriver
+import numpy as np
+from mpi4py import MPI
+
+class OnewayStructTrimDriver(OnewayStructDriver):
+
+    """
+    goal of this class is to run oneway-coupled sizing
+    while also trimming the wing using a set of loads obtained 
+    when pull up might not be satisfied..
+
+    aero / struct loads are scaled up as an AOA variable is scaled up
+    for all uncoupled scenarios (have scenario.coupled = False)
+
+    the user should setup a load factor based composite function
+    such as lift - load_factor * weight = 0
+    where load_factor is user-specified for that scenario
+    """
+
+    def __init__(
+        self,
+        solvers,
+        model,
+        initial_trim_dict:dict,
+        transfer_settings=None,
+        nprocs=None,
+        fun3d_dir=None,
+        external_shape=False,
+        timing_file=None,
+    ):
+
+        # create base class OnewayStructDriver
+        super(OnewayStructTrimDriver,self).__init__(
+            solvers,
+            model,
+            transfer_settings,
+            nprocs,
+            fun3d_dir,
+            external_shape,
+            timing_file,
+        )
+
+        # get data from scenario initial trim dict
+        # assumed to hold initial values for (not case sensitive)
+        # and lift is C_L normalized by area and qinf
+        # {scenario_name : {'cl' : cl_0, 'AOA' : AOA_0}}
+        # only required to put scenarios which are uncoupled here
+        # with scenario.coupled = False boolean
+
+        self.initial_trim_dict = initial_trim_dict
+
+        # save initial struct loads vectors for each scenario
+        self._orig_struct_loads = {}
+        self.uncoupled_scenarios = [scenario for scenario in model.scenarios if not(scenario.coupled)]
+        for scenario in self.uncoupled_scenarios:
+            self._orig_struct_loads[scenario.name] = {}
+            for body in model.bodies:
+                struct_loads = body.struct_loads[scenario.id]
+                self._orig_struct_loads[scenario.name][body.name] = struct_loads * 1.0
+
+    @classmethod
+    def prime_loads_from_file(
+        cls,
+        filename,
+        solvers,
+        initial_trim_dict,
+        model,
+        nprocs,
+        transfer_settings,
+        external_shape=False,
+        init_transfer=False,
+        timing_file=None,
+    ):
+        # same as base class prime_loads_from_file but with extra input argument
+        # aka initial_trim_dict
+        comm = solvers.comm
+        world_rank = comm.Get_rank()
+        if world_rank < nprocs:
+            color = 1
+        else:
+            color = MPI.UNDEFINED
+        tacs_comm = comm.Split(color, world_rank)
+
+        # initialize transfer settings
+        comm_manager = solvers.comm_manager
+
+        # read in the loads from the file
+        loads_data = model._read_aero_loads(comm, filename)
+
+        # initialize the transfer scheme then distribute aero loads
+        for body in model.bodies:
+            body.initialize_transfer(
+                comm=comm,
+                struct_comm=tacs_comm,
+                struct_root=comm_manager.struct_root,
+                aero_comm=comm_manager.aero_comm,
+                aero_root=comm_manager.aero_root,
+                transfer_settings=transfer_settings,
+            )
+            for scenario in model.scenarios:
+                body.initialize_variables(scenario)
+                assert scenario.steady
+            body._distribute_aero_loads(loads_data, steady=True)
+
+        tacs_driver = cls(
+            solvers,
+            model,
+            initial_trim_dict,
+            nprocs=nprocs,
+            external_shape=external_shape,
+            timing_file=timing_file,
+        )
+        if init_transfer:
+            tacs_driver._transfer_fixed_aero_loads()
+
+        return tacs_driver
+
+    def solve_forward(self):
+
+        # scale up the loads by new AOA vs previous AOA
+        # note this only works for steady-state case
+        for scenario in self.uncoupled_scenarios:
+            orig_AOA = self.initial_trim_dict[scenario.name]['AOA']
+            new_AOA = scenario.get_variable('AOA').value.real
+            for body in self.model.bodies:
+                orig_struct_loads = self._orig_struct_loads[scenario.name][body.name]
+                body.struct_loads[scenario.id][:] = (orig_struct_loads * new_AOA / orig_AOA)[:]
+
+        # now do super class solve_forward which will include
+        # transferring fixed aero loads to the new struct loads and then linear static solve
+        super(OnewayStructTrimDriver,self).solve_forward()
+
+        # compute new lift values, for function name cl
+        for scenario in self.uncoupled_scenarios:
+            orig_cl = self.initial_trim_dict[scenario.name]['cl']
+            orig_AOA = self.initial_trim_dict[scenario.name]['AOA']
+            new_AOA = scenario.get_variable('AOA').value.real
+
+            for func in scenario.functions:
+                if func.name == 'cl':
+                    func.value = orig_cl * new_AOA / orig_AOA
+
+        # composite functions are evaluated in the OptimizationManager FYI and will also be updated after this..
+
+    def solve_adjoint(self):
+
+        # do super class solve_adjoint (same adjoint solve as before)
+        # since modified f_A is output of adjoint solve and not coupled...
+        # so doesn't matter really
+        super(OnewayStructTrimDriver,self).solve_adjoint()
+
+    def _solve_steady_adjoint(self, scenario, bodies):
+        super()._solve_steady_adjoint(scenario, bodies)
+
+        # get additional derivative terms for custom
+        self._get_custom_derivatives(scenario)
+
+    def _solve_unsteady_adjoint(self, scenario, bodies):
+        super()._solve_unsteady_adjoint(scenario, bodies)
+
+        # get additional derivative terms for custom
+        self._get_custom_derivatives(scenario)
+
+    def _get_custom_derivatives(self, scenario):
+        """get custom trim derivatives, this is used in the """
+
+        orig_cl = self.initial_trim_dict[scenario.name]['cl']
+        orig_AOA = self.initial_trim_dict[scenario.name]['AOA']
+        aoa_var = scenario.get_variable('AOA')
+
+        # since mass not adjoint function only iterate over these guys
+        adjoint_functions = [func for func in scenario.functions if func.adjoint]
+        for ifunc,func in enumerate(adjoint_functions):
+            if func.name == 'cl':
+                func.derivatives[aoa_var] = orig_cl / orig_AOA 
+                continue
+
+            # account for changing loads terms in AOA
+            AOA_deriv = 0.0
+            for body in self.model.bodies:
+                struct_loads_ajp = body.get_struct_loads_ajp(scenario)
+                func_fs_ajp = struct_loads_ajp[:,ifunc]
+                orig_struct_loads = self._orig_struct_loads[scenario.name][body.name]
+                free_struct_loads = orig_struct_loads * 1.0
+
+                # this didn't change anything
+                # # temp - try to zero BCs at ext force locations?
+                # structural = self.solvers.structural
+                # assembler = structural.assembler
+                # ext_force = structural.ext_force
+                # ext_force_array = ext_force.getArray()
+                # ndof = assembler.getVarsPerNode()
+                # for i in range(3): 
+                #     # set only 0,1,2 into ext_force, then we will apply BCs to zero out dirichlet BCs
+                #     ext_force_array[i::ndof] = free_struct_loads[i::3]
+                # assembler.setBCs(ext_force) # zero out forces at dirichlet BCs (since have no effect on structure)
+                # for i in range(3):
+                #     free_struct_loads[i::3] = ext_force_array[i::ndof]         
+
+                AOA_deriv += np.dot(func_fs_ajp, free_struct_loads / orig_AOA)
+
+            # add across all processors then reduce
+            global_derivative = self.comm.reduce(AOA_deriv, op=MPI.SUM, root=0)
+            global_derivative = self.comm.bcast(global_derivative)
+
+            func.derivatives[aoa_var] = global_derivative

funtofem/driver/funtofem_shape_driver.py

@@ -215,7 +215,7 @@ def __init__(
                     self._flow_solver_type = "fun3d"
             # TBD on new types
         else:  # check with shape change
-            if fun3d_loader is not None:
+            if fun3d_loader is not None and caps_loader is not None:
                 if isinstance(model.flow, Fun3dModel):
                     self._flow_solver_type = "fun3d"
                     self.flow_aim = model.flow.fun3d_aim