Update optimizer.py

modules/optimizer.py

@@ -34,7 +34,6 @@ class Optimizer(object):
         self.to(self.device)
 
         
-
     def to(self, device):
         self.device = device
         self.surrogate_model.to(device)
@@ -44,25 +43,78 @@ class Optimizer(object):
         # keep the size of the detector within 3m
         raw_detector_parameters = self.detector_parameters 
         # this creates a loss that increases starting from 3 meters using relu
-        detector_length = torch.sum(raw_detector_parameters)
+        detector_length = torch.sum(raw_detector_parameters[:6])# changed this here so as to not include material choice in this
+
+        thickness_parameters = self.detector_parameters[0:6]
+        material_parameters = self.detector_parameters[6:12]
+        abs_parameters = self.detector_parameters[6:9]
+        scint_parameters = self.detector_parameters[9:12]
+
+        # cost_dic = {"G4_POLYSTYRENE": 0.001
+        #     ,"G4_PLASTIC_SC_VINYLTOLUENE": 0.001
+        #     ,"G4_BGO": 666.66
+        #     ,"G4_LSO": 5000
+        #     ,"G4_LYSO": 75000
+        #     ,"G4_CESIUM_IODIDE": 0.001
+        #     ,"G4_Pb": 25
+        #     ,"G4_Fe": 4.166
+        #     ,"G4_W": 500
+        #     ,"G4_Cu": 66.66
+        #     ,"G4_BRASS": 16.66
+        #     ,"G4_Si": 8330000
+        #     ,"G4_PbWO4": 2500 }
+        
+
+        
         if self.constraints is not None:
             if 'length' in self.constraints:
-                detector_length = torch.sum(raw_detector_parameters)
+                detector_length = torch.sum(raw_detector_parameters[:6])# changed this here so as to not include material choice in this
                 total_length_loss = torch.mean(100.*torch.nn.ReLU()(detector_length - self.constraints['length'])**2)#constrain it to 25cm
+            if 'lower' in self.constraints:
+                lbound = (torch.ones(12) * self.constraints['lower']).to(self.device)
+                lower_loss = torch.mean(1000.*torch.nn.ReLU()(lbound- raw_detector_parameters)**2)
+            if 'upper' in self.constraints:
+                ubound = (torch.ones(6) * self.constraints['upper']).to(self.device)
+                upper_loss = torch.mean(100.*torch.nn.ReLU()(material_parameters - ubound)**2)
+            if 'diff' in self.constraints:
+                abs_maxdiff = max(abs_parameters) - min(abs_parameters)
+                scint_maxdiff = max(scint_parameters)- min(scint_parameters)
+
+                diff_loss = torch.mean(100.*torch.nn.ReLU()(abs_maxdiff - self.constraints['diff'])**2)
+                diff_loss += torch.mean(100.*torch.nn.ReLU()(scint_maxdiff - self.constraints['diff'])**2)
+            if 'cost' in self.constraints: 
+
+                abs_mask = (abs_parameters >= 0.5)
+                cost_abs = torch.where(abs_mask, 25 , 4.166)
+                scint_mask = (scint_parameters >= 0.5)
+                cost_scint = torch.where(scint_mask, 2500, 0.001)
+
+                combined = torch.cat((cost_abs, cost_scint), dim=0)
+
+                combined_cost = combined.to(self.device)
+
+                # bounded_thickness = torch.tensor(thickness_parameters)
+                # bounded_thickness = torch.where(bounded_thickness > 0, bounded_thickness, 0)
+
+                cost = torch.sum(combined_cost * thickness_parameters)
+                cost_loss = torch.mean(50.* torch.nn.ReLU()(cost - self.constraints['cost'])**2)
+
+        
+                             
         # now keep parameters such that within the box size of the generator, there are always some positive values even if the 
         # central parameters are negative. Both box size and raw_detector_parameters are in non-normalised space, so this is straight forward
         # the generator will have to provide the box size
         # this will avoid mode collapse
 
-        box = self.cu_box # the box has same ordering as raw_detector_parameters
-        lower_para_bound = -box/1.1
-        bloss = torch.mean(100.*torch.nn.ReLU()(lower_para_bound - raw_detector_parameters)**2)
+        # box = self.cu_box # the box has same ordering as raw_detector_parameters
+        # lower_para_bound = -box/1.1
+        # bloss = torch.mean(100.*torch.nn.ReLU()(lower_para_bound - raw_detector_parameters)**2)
+        # + bloss 
 
-        return total_length_loss + bloss
+        return total_length_loss +lower_loss + upper_loss + diff_loss + cost_loss
     
     def clamp_parameters(self):
-        return #for now
-    
+        return
         self.detector_parameters = self.detector_parameters.clamp(1e-3) #DEBUG, NE
 
     def adjust_generator_covariance(self, direction, min_scale=2.0):