Add design variable support for aux elements

src/TACSAssembler.cpp

@@ -160,6 +160,8 @@ TACSAssembler::TACSAssembler(MPI_Comm _tacs_comm, int _varsPerNode,
   designNodeMap = NULL;
   designExtDist = NULL;
   designDepNodes = NULL;
+  numGlobalDVs = 0;
+  globalDVNums = NULL;
 
   // Set the local element data to NULL
   elementData = NULL;
@@ -281,6 +283,9 @@ TACSAssembler::~TACSAssembler() {
   if (elementSensIData) {
     delete[] elementSensIData;
   }
+  if (globalDVNums) {
+    delete[] globalDVNums;
+  }
 
   // Delete initial condition vectors
   if (vars0) {
@@ -972,6 +977,7 @@ void TACSAssembler::setAuxElements(TACSAuxElements *auxElems) {
     auxElements->decref();
   }
   auxElements = auxElems;
+  bool maxDVsUpdated = false;
 
   // Check whether the auxiliary elements match
   if (auxElements) {
@@ -987,6 +993,23 @@ void TACSAssembler::setAuxElements(TACSAuxElements *auxElems) {
                  aux[k].elem->getNumVariables()) {
         fprintf(stderr, "[%d] Auxiliary element sizes do not match\n", mpiRank);
       }
+
+      // Update the maximum number of element design variables
+      int numDVs = aux[k].elem->getDesignVarNums(aux[k].num, 0, NULL);
+      if (numDVs > maxElementDesignVars) {
+        maxElementDesignVars = numDVs;
+        maxDVsUpdated = true;
+      }
+    }
+
+    // If the buffers have already been allocated and the new max exceeds their
+    // current size, reallocate them so downstream code doesn't overflow.
+    if (maxDVsUpdated && elementSensData && elementSensIData) {
+      delete[] elementSensData;
+      delete[] elementSensIData;
+      elementSensData =
+          new TacsScalar[designVarsPerNode * maxElementDesignVars];
+      elementSensIData = new int[maxElementDesignVars];
     }
   }
 }
@@ -2617,9 +2640,18 @@ int TACSAssembler::initialize() {
 
   // Create the design variable node mapping
   if (!designNodeMap) {
-    // Get the number of design variables
+    // Get the number of design variables from element scan
     int numDVs = getNumDesignVars();
 
+    // Extend to cover any global DV indices not referenced by regular elements
+    // (e.g. DVs that will only be used by aux elements added after
+    // initialize())
+    for (int i = 0; i < numGlobalDVs; i++) {
+      if (globalDVNums[i] + 1 > numDVs) {
+        numDVs = globalDVNums[i] + 1;
+      }
+    }
+
     if (mpiRank > 0) {
       numDVs = 0;
     }
@@ -2672,6 +2704,13 @@ int TACSAssembler::initialize() {
       }
     }
 
+    // Count Python-level global DV indices not locally owned
+    for (int i = 0; i < numGlobalDVs; i++) {
+      if (globalDVNums[i] < lower || globalDVNums[i] >= upper) {
+        dvLen++;
+      }
+    }
+
     // Allocate space for absolutely everything!
     int *allDVs = new int[dvLen];
 
@@ -2711,6 +2750,14 @@ int TACSAssembler::initialize() {
       }
     }
 
+    // Add any global DV indices not locally owned
+    for (int i = 0; i < numGlobalDVs; i++) {
+      if (globalDVNums[i] < lower || globalDVNums[i] >= upper) {
+        allDVs[dvLen] = globalDVNums[i];
+        dvLen++;
+      }
+    }
+
     dvLen = TacsUniqueSort(dvLen, allDVs);
     int *dvs = new int[dvLen];
     memcpy(dvs, allDVs, dvLen * sizeof(int));
@@ -3062,6 +3109,39 @@ void TACSAssembler::setDesignNodeMap(int _designVarsPerNode,
   designNodeMap = _designNodeMap;
 }
 
+/**
+  Store the global DV indices that must be distributed to all procs.
+
+  These indices are appended to the external DV distribution during
+  initialize(), ensuring that aux elements added after initialize() can
+  always access them via dvs->getValues().
+
+  @param n Number of global DV indices
+  @param dvNums Array of global DV indices (length n)
+*/
+void TACSAssembler::setGlobalDVIndices(int n, const int *dvNums) {
+  delete[] globalDVNums;
+  numGlobalDVs = n;
+  globalDVNums = new int[n];
+  memcpy(globalDVNums, dvNums, n * sizeof(int));
+}
+
+/**
+  Retrieve the global DV indices previously registered via setGlobalDVIndices().
+
+  Sets *dvNums to point to the internal array (do not free or modify).
+
+  @param dvNums Output pointer to the internal array of global DV indices
+  @return The number of global DV indices (same as n passed to
+  setGlobalDVIndices)
+*/
+int TACSAssembler::getGlobalDVIndices(const int **dvNums) {
+  if (dvNums) {
+    *dvNums = globalDVNums;
+  }
+  return numGlobalDVs;
+}
+
 /**
   Set the dependent design variable information
 

src/TACSAssembler.h

@@ -100,6 +100,12 @@ class TACSAssembler : public TACSObject {
   // --------------------------------------------------
   void setDesignNodeMap(int _designVarsPerNode,
                         TACSNodeMap *_designVarMap = NULL);
+
+  // Register global DV indices that must be distributed to all procs.
+  // These are added to designExtDist during initialize() so that aux
+  // elements added after initialize() can always access them.
+  void setGlobalDVIndices(int n, const int *dvNums);
+  int getGlobalDVIndices(const int **dvNums);
   int setDesignDependentNodes(int numDepDesignVars, const int *_depNodePtr,
                               const int *_depNodes,
                               const double *_depNodeWeights);
@@ -406,6 +412,8 @@ class TACSAssembler : public TACSObject {
   int numDependentNodes;   // number of dependent nodes
   int numMultiplierNodes;  // number of multiplier nodes/elements
   int designVarsPerNode;   // number of design variables at each design "node"
+  int numGlobalDVs;        // number of Python-level global DV indices
+  int *globalDVNums;       // the global DV indices
 
   // Maximum element information
   int maxElementDesignVars;  // maximum number of design variable

src/TACSCreator.cpp

@@ -99,6 +99,10 @@ TACSCreator::TACSCreator(MPI_Comm _comm, int _vars_per_node) {
   owned_elements = NULL;
   owned_nodes = NULL;
 
+  // No global DVs registered yet
+  numGlobalDVs = 0;
+  globalDVNums = NULL;
+
   // Set the elements array to NULL
   elements = NULL;
   element_creator = NULL;
@@ -162,6 +166,10 @@ TACSCreator::~TACSCreator() {
     delete[] local_elem_id_nums;
   }
 
+  if (globalDVNums) {
+    delete[] globalDVNums;
+  }
+
   if (elements) {
     for (int i = 0; i < num_elem_ids; i++) {
       if (elements[i]) {
@@ -291,6 +299,16 @@ void TACSCreator::setElementCreator(TACSElement *(*func)(int, int)) {
   element_creator = func;
 }
 
+/*
+  Store the global DV indices so they get included in designExtDist on all procs
+*/
+void TACSCreator::setGlobalDVIndices(int n, const int *dvNums) {
+  delete[] globalDVNums;
+  numGlobalDVs = n;
+  globalDVNums = new int[n];
+  memcpy(globalDVNums, dvNums, n * sizeof(int));
+}
+
 /*
   Set the nodal locations
 */
@@ -873,6 +891,11 @@ TACSAssembler *TACSCreator::createTACS() {
   delete[] bc_vals;
   bc_vals = NULL;
 
+  // Pass global DV indices so they appear in designExtDist on all procs
+  if (numGlobalDVs > 0) {
+    tacs->setGlobalDVIndices(numGlobalDVs, globalDVNums);
+  }
+
   // Use the reordering if the flag has been set in the
   // TACSCreator object
   if (use_reordering) {

src/TACSCreator.h

@@ -79,6 +79,10 @@ class TACSCreator : public TACSObject {
   // ------------------
   void partitionMesh(int split_size = 0, const int *part = NULL);
 
+  // Set global DV indices to distribute to all procs
+  // -------------------------------------------------
+  void setGlobalDVIndices(int n, const int *dvNums);
+
   // Set the elements into TACS creator
   // ----------------------------------
   void setElements(int _num_elems, TACSElement **_elements);
@@ -143,6 +147,10 @@ class TACSCreator : public TACSObject {
   // The node locations
   TacsScalar *Xpts;
 
+  // Global DV indices to force into designExtDist on all procs
+  int numGlobalDVs;
+  int *globalDVNums;
+
   // Elements and the corresponding element id numbers
   int num_elem_ids;
   TACSElement **elements;

src/elements/TACSConvectiveTraction2D.cpp

@@ -86,7 +86,7 @@ void TACSConvectiveTraction2D::addResidual(
     double weight = basis->getFaceQuadraturePoint(faceIndex, n, pt, tangent);
 
     // Get the face normal
-    TacsScalar X[3], Xd[4], normal[3];
+    TacsScalar X[3], Xd[6], normal[3];
     TacsScalar area = basis->getFaceNormal(faceIndex, n, Xpts, X, Xd, normal);
 
     // Get the field value
@@ -132,7 +132,7 @@ void TACSConvectiveTraction2D::addJacobian(
     double weight = basis->getFaceQuadraturePoint(faceIndex, n, pt, tangent);
 
     // Get the face normal
-    TacsScalar X[3], Xd[4], normal[3];
+    TacsScalar X[3], Xd[6], normal[3];
     TacsScalar area = basis->getFaceNormal(faceIndex, n, Xpts, X, Xd, normal);
 
     // Get the field value

src/elements/TACSElement.h

@@ -203,7 +203,8 @@ class TACSElement : public TACSObject {
     Possibly NULL.
   */
   virtual TACSElement *createElementTraction(int faceIndex,
-                                             const TacsScalar t[]) {
+                                             const TacsScalar t[],
+                                             const int *tracDVNums = NULL) {
     return NULL;
   }
 
@@ -213,7 +214,8 @@ class TACSElement : public TACSObject {
     @return The TACSElement pressure class associated with this element.
     Possibly NULL.
   */
-  virtual TACSElement *createElementPressure(int faceIndex, TacsScalar p) {
+  virtual TACSElement *createElementPressure(int faceIndex, TacsScalar p,
+                                             int pressureDVNum = -1) {
     return NULL;
   }
 
@@ -223,7 +225,8 @@ class TACSElement : public TACSObject {
     @return The TACSElement inertial force class associated with this element.
     Possibly NULL.
   */
-  virtual TACSElement *createElementInertialForce(const TacsScalar g[]) {
+  virtual TACSElement *createElementInertialForce(
+      const TacsScalar g[], const int *inertiaVecDVNums = NULL) {
     return NULL;
   }
 

src/elements/TACSElement2D.cpp

@@ -53,21 +53,24 @@ TACSElementBasis *TACSElement2D::getElementBasis() { return basis; }
 TACSElementModel *TACSElement2D::getElementModel() { return model; }
 
 TACSElement *TACSElement2D::createElementTraction(int faceIndex,
-                                                  const TacsScalar t[]) {
+                                                  const TacsScalar t[],
+                                                  const int *tracDVNums) {
   int varsPerNode = getVarsPerNode();
-  return new TACSTraction2D(varsPerNode, faceIndex, basis, t);
+  return new TACSTraction2D(varsPerNode, faceIndex, basis, t, 1, tracDVNums);
 }
 
-TACSElement *TACSElement2D::createElementPressure(int faceIndex, TacsScalar p) {
+TACSElement *TACSElement2D::createElementPressure(int faceIndex, TacsScalar p,
+                                                  int pressureDVNum) {
   int varsPerNode = getVarsPerNode();
-  return new TACSPressure2D(varsPerNode, faceIndex, basis, p);
+  return new TACSPressure2D(varsPerNode, faceIndex, basis, p, pressureDVNum);
 }
 
 TACSElement *TACSElement2D::createElementInertialForce(
-    const TacsScalar inertiaVec[]) {
+    const TacsScalar inertiaVec[], const int *inertiaVecDVNums) {
   int varsPerNode = getVarsPerNode();
   TACSConstitutive *con = model->getConstitutive();
-  return new TACSInertialForce2D(varsPerNode, con, basis, inertiaVec);
+  return new TACSInertialForce2D(varsPerNode, con, basis, inertiaVec,
+                                 inertiaVecDVNums);
 }
 
 int TACSElement2D::getNumQuadraturePoints() {

src/elements/TACSElement2D.h

@@ -34,9 +34,12 @@ class TACSElement2D : public TACSElement {
   ElementType getElementType();
   TACSElementBasis *getElementBasis();
   TACSElementModel *getElementModel();
-  TACSElement *createElementTraction(int faceIndex, const TacsScalar t[]);
-  TACSElement *createElementPressure(int faceIndex, TacsScalar p);
-  TACSElement *createElementInertialForce(const TacsScalar inertiaVec[]);
+  TACSElement *createElementTraction(int faceIndex, const TacsScalar t[],
+                                     const int *tracDVNums = NULL);
+  TACSElement *createElementPressure(int faceIndex, TacsScalar p,
+                                     int pressureDVNum = -1);
+  TACSElement *createElementInertialForce(const TacsScalar inertiaVec[],
+                                          const int *inertiaVecDVNums = NULL);
   int getNumQuadraturePoints();
   double getQuadratureWeight(int n);
   double getQuadraturePoint(int n, double pt[]);

src/elements/TACSElement3D.cpp

@@ -46,21 +46,24 @@ int TACSElement3D::getDesignVarsPerNode() {
 }
 
 TACSElement *TACSElement3D::createElementTraction(int faceIndex,
-                                                  const TacsScalar t[]) {
+                                                  const TacsScalar t[],
+                                                  const int *tracDVNums) {
   int varsPerNode = getVarsPerNode();
-  return new TACSTraction3D(varsPerNode, faceIndex, basis, t);
+  return new TACSTraction3D(varsPerNode, faceIndex, basis, t, 1, tracDVNums);
 }
 
-TACSElement *TACSElement3D::createElementPressure(int faceIndex, TacsScalar p) {
+TACSElement *TACSElement3D::createElementPressure(int faceIndex, TacsScalar p,
+                                                  int pressureDVNum) {
   int varsPerNode = getVarsPerNode();
-  return new TACSPressure3D(varsPerNode, faceIndex, basis, p);
+  return new TACSPressure3D(varsPerNode, faceIndex, basis, p, pressureDVNum);
 }
 
 TACSElement *TACSElement3D::createElementInertialForce(
-    const TacsScalar inertiaVec[]) {
+    const TacsScalar inertiaVec[], const int *inertiaVecDVNums) {
   int varsPerNode = getVarsPerNode();
   TACSConstitutive *con = model->getConstitutive();
-  return new TACSInertialForce3D(varsPerNode, con, basis, inertiaVec);
+  return new TACSInertialForce3D(varsPerNode, con, basis, inertiaVec,
+                                 inertiaVecDVNums);
 }
 
 TACSElement *TACSElement3D::createElementCentrifugalForce(

src/elements/TACSElement3D.h

@@ -35,9 +35,12 @@ class TACSElement3D : public TACSElement {
   ElementType getElementType();
   TACSElementBasis *getElementBasis();
   TACSElementModel *getElementModel();
-  TACSElement *createElementTraction(int faceIndex, const TacsScalar t[]);
-  TACSElement *createElementPressure(int faceIndex, TacsScalar p);
-  TACSElement *createElementInertialForce(const TacsScalar inertiaVec[]);
+  TACSElement *createElementTraction(int faceIndex, const TacsScalar t[],
+                                     const int *tracDVNums = NULL);
+  TACSElement *createElementPressure(int faceIndex, TacsScalar p,
+                                     int pressureDVNum = -1);
+  TACSElement *createElementInertialForce(const TacsScalar inertiaVec[],
+                                          const int *inertiaVecDVNums = NULL);
   TACSElement *createElementCentrifugalForce(const TacsScalar omegaVec[],
                                              const TacsScalar rotCenter[],
                                              const bool first_order = false);