fix tests, formatting

Author: dbochkov-flexcompute

tests/test_components/test_heat.py

@@ -11,6 +11,7 @@
     HeatFluxBC,
     HeatSimulation,
     HeatSimulationData,
+    HeatSource,
     MediumMediumInterface,
     SimulationBoundary,
     SolidSpec,
@@ -20,7 +21,6 @@
     TemperatureBC,
     TemperatureData,
     TemperatureMonitor,
-    UniformHeatSource,
     UniformUnstructuredGrid,
 )
 from tidy3d.exceptions import DataError
@@ -251,20 +251,30 @@ def test_grid_spec():
 
 
 def make_heat_source():
-    return UniformHeatSource(structures=["solid_structure"], rate=100)
+    return HeatSource(structures=["solid_structure"], rate=100)
+
+
+def make_custom_heat_source():
+    return HeatSource(
+        structures=["solid_structure"],
+        rate=td.SpatialDataArray(np.ones((1, 2, 3)), coords=dict(x=[0], y=[1, 2], z=[3, 4, 5])),
+    )
 
 
 def test_heat_source():
     source = make_heat_source()
+    source = make_custom_heat_source()
     with pytest.raises(pd.ValidationError):
         _ = source.updated_copy(structures=[])
 
 
-def make_heat_sim():
+def make_heat_sim(include_custom_source: bool = True):
     fluid_medium, solid_medium = make_heat_mediums()
     fluid_structure, solid_structure = make_heat_structures()
     bc_temp, bc_flux, bc_conv = make_heat_bcs()
-    heat_source = make_heat_source()
+    sources = [make_heat_source()]
+    if include_custom_source:
+        sources += [make_custom_heat_source()]
 
     pl1 = HeatBoundarySpec(
         condition=bc_conv, placement=MediumMediumInterface(mediums=["fluid_medium", "solid_medium"])
@@ -292,7 +302,7 @@ def make_heat_sim():
         size=(2, 2, 2),
         boundary_spec=[pl1, pl2, pl3, pl4, pl5],
         grid_spec=grid_spec,
-        sources=[heat_source],
+        sources=sources,
         monitors=temp_mnts,
     )
 
@@ -323,7 +333,7 @@ def test_heat_sim():
             _ = heat_sim.updated_copy(boundary_spec=[pl])
 
     with pytest.raises(pd.ValidationError):
-        _ = heat_sim.updated_copy(sources=[UniformHeatSource(structures=["noname"])], rate=-10)
+        _ = heat_sim.updated_copy(sources=[HeatSource(structures=["noname"])], rate=-10)
 
     # run 2D case
     _ = heat_sim.updated_copy(center=(0.7, 0, 0), size=(0, 2, 2), monitors=heat_sim.monitors[:5])

tests/test_web/test_webapi_heat.py

@@ -259,7 +259,7 @@ def test_estimate_cost(set_api_key, mock_get_info, mock_metadata):
 
 @responses.activate
 def test_download_json(monkeypatch, mock_get_info, tmp_path):
-    sim = make_heat_sim()
+    sim = make_heat_sim(include_custom_source=False)
 
     def mock_download(*args, **kwargs):
         pass
@@ -278,7 +278,7 @@ def get_str(*args, **kwargs):
 @responses.activate
 def test_load_simulation(monkeypatch, mock_get_info, tmp_path):
     def mock_download(*args, **kwargs):
-        make_heat_sim().to_file(args[1])
+        make_heat_sim(include_custom_source=False).to_file(args[1])
 
     monkeypatch.setattr(f"{task_core_path}.SimulationTask.get_simulation_json", mock_download)
 

tidy3d/components/tcad/simulation/heat_charge.py

@@ -1494,28 +1494,16 @@ def _add_source_cbar(self, ax: Ax, property: str = "heat_conductivity"):
             ax=ax,
         )
 
-    def _safe_float_conversion(self, string) -> float:
-        """Function to deal with failed string2float conversion when using
-        expressions in 'HeatSource'"""
-        try:
-            return float(string)
-        except ValueError:
-            return None
-
     def source_bounds(self, property: str = "heat_conductivity") -> Tuple[float, float]:
         """Compute range of heat sources present in the simulation."""
 
         if property == "heat_conductivity" or property == "source":
             rate_list = [
-                self._safe_float_conversion(source.rate)
-                for source in self.sources
-                if isinstance(source, HeatSource)
+                np.mean(source.rate) for source in self.sources if isinstance(source, HeatSource)
             ]
         elif property == "electric_conductivity":
             rate_list = [
-                self._safe_float_conversion(source.rate)
-                for source in self.sources
-                if isinstance(source, ChargeSourceTypes)
+                source.rate for source in self.sources if isinstance(source, ChargeSourceTypes)
             ]  # this is currently an empty list
 
         rate_list.append(0)
@@ -1537,7 +1525,7 @@ def _get_structure_source_plot_params(
             plot_params = plot_params.copy(update={"alpha": alpha})
 
         if isinstance(source, HeatSource):
-            rate = self._safe_float_conversion(source.rate)
+            rate = np.mean(source.rate)
             if rate is not None:
                 delta_rate = rate - source_min
                 delta_rate_max = source_max - source_min + 1e-5

tidy3d/components/tcad/source/heat.py

@@ -6,9 +6,9 @@
 
 import pydantic.v1 as pd
 
+from tidy3d.components.data.data_array import SpatialDataArray
 from tidy3d.components.tcad.source.abstract import StructureBasedHeatChargeSource
 from tidy3d.constants import VOLUMETRIC_HEAT_RATE
-from tidy3d.components.data.data_array import SpatialDataArray
 from tidy3d.log import log