Merge pull request atmtools#225 from lkluft/add-deepcopy

Add deepcopy

howto/atmosphere.md

@@ -53,7 +53,7 @@ Finally, we can plot the equilibrated temperature profile
 
 ```{code-cell} ipython3
 fig, ax = plt.subplots()
-plots.profile_p_log(atmosphere['plev'], atmosphere['T'][-1, :])
+plots.profile_p_log(rce.atmosphere['plev'], rce.atmosphere['T'][-1, :])
 ax.set_xlabel(r"$T$ / K")
 ax.set_ylabel("$p$ / hPa")
 ```

howto/clouds.md

@@ -43,12 +43,12 @@ rce = konrad.RCE(
 rce.run()
 
 fig, (ax0, ax1) = plt.subplots(ncols=2, sharey=True)
-plots.profile_p_log(atmosphere["plev"], atmosphere["T"][-1], ax=ax0)
+plots.profile_p_log(rce.atmosphere["plev"], rce.atmosphere["T"][-1], ax=ax0)
 ax0.set_ylabel("$p$ / hPa")
 ax0.set_xlabel("$T$ / K")
 
 ax1.axvline(0, color="k", linewidth=0.8)
-plots.profile_p_log(atmosphere["plev"], rce.radiation["net_htngrt"][-1], ax=ax1)
+plots.profile_p_log(rce.atmosphere["plev"], rce.radiation["net_htngrt"][-1], ax=ax1)
 ax1.set_xlabel("Q / $\sf K\,day^{-1}$")
 ax1.set_xlim(-4, 0.5)
 ax1.set_ylim(bottom=phlev.max())
@@ -58,7 +58,7 @@ ax1.set_ylim(bottom=phlev.max())
 
 ```{code-cell} ipython3
 single_cloud = konrad.cloud.ConceptualCloud(
-    atmosphere,  # required for consistent coordinates
+    rce.atmosphere,  # required for consistent coordinates
     cloud_top=200e2,  # in Pa
     depth=100e2,  # in Pa
     phase="ice",  # "ice" or "liquid"
@@ -68,20 +68,20 @@ single_cloud = konrad.cloud.ConceptualCloud(
 
 rrtmg = konrad.radiation.RRTMG()
 rrtmg.update_heatingrates(
-    atmosphere=atmosphere,
+    atmosphere=rce.atmosphere,
     cloud=single_cloud,
 )
 
 fig, (ax0, ax1) = plt.subplots(ncols=2, sharey=True)
 ax0.axvline(0, color="k", linewidth=0.8)
-plots.profile_p_log(atmosphere["plev"], rrtmg["net_htngrt"][-1], ax=ax0)
+plots.profile_p_log(rce.atmosphere["plev"], rrtmg["net_htngrt"][-1], ax=ax0)
 ax0.set_xlabel("Q / $\sf K\,day^{-1}$")
 ax0.set_xlim(-4, 0.5)
 ax0.set_ylabel("$p$ / hPa")
 ax0.set_ylim(bottom=phlev.max())
 
 ax1.axvline(0, color="k", linewidth=0.8)
-plots.profile_p_log(atmosphere["plev"], rrtmg["net_htngrt"][-1] - rrtmg["net_htngrt_clr"][-1], ax=ax1)
+plots.profile_p_log(rce.atmosphere["plev"], rrtmg["net_htngrt"][-1] - rrtmg["net_htngrt_clr"][-1], ax=ax1)
 ax1.set_xlabel("$\sf Q_\mathrm{cloud}$ / $\sf K\,day^{-1}$")
 ax1.set_xlim(-2.25, 2.25)
 ax1.set_ylim(bottom=phlev.max())

howto/convection.md

@@ -53,7 +53,7 @@ In a second step, we enable the convective adjustment. We copy the existing atmo
 
 ```{code-cell} ipython3
 rce = konrad.RCE(
-    atmosphere.copy(),  # Create an separate atmosphere component.
+    atmosphere,
     convection=konrad.convection.HardAdjustment(),
     timestep='24h',
     max_duration='150d',

howto/feedback.md

@@ -45,10 +45,10 @@ spinup = konrad.RCE(
 )
 spinup.run()  # Start the simulation.
 
-atmosphere["CO2"][:] *= 2.0 # double the CO2 concentration
+spinup.atmosphere["CO2"][:] *= 2.0 # double the CO2 concentration
 
 perturbed = konrad.RCE(
-    atmosphere,
+    spinup.atmosphere,
     surface=konrad.surface.SlabOcean(
         temperature=295.0,
         heat_sink=spinup.radiation["toa"][-1],

howto/forcing.md

@@ -54,8 +54,8 @@ atmospheric state, especially the temperature, occur.
 
 ```{code-cell} ipython3
 # Calculate OLR at perturbed atmospheric state.
-atmosphere["CO2"][:] *= 2  # double the CO2 concentration
-spinup.radiation.update_heatingrates(atmosphere)
+spinup.atmosphere["CO2"][:] *= 2  # double the CO2 concentration
+spinup.radiation.update_heatingrates(spinup.atmosphere)
 
 instant_forcing = -(spinup.radiation["lw_flxu"][-1] - olr_ref)
 ```
@@ -94,7 +94,7 @@ The effective forcing includes the so called "stratospheric adjustment".
 One can simulated by keeping the surface temperature fixed but allowing the atmopsheric temperature to adjust to the increased CO2 concentration.
 
 ```{code-cell} ipython3
-perturbed = konrad.RCE(atmosphere.copy(), timestep='24h',max_duration='150d')
+perturbed = konrad.RCE(spinup.atmosphere, timestep='24h',max_duration='150d')
 perturbed.run()
 
 effective_forcing = -(perturbed.radiation["lw_flxu"][-1] - olr_ref)

howto/getting_started.md

@@ -68,7 +68,9 @@ Finally, we can plot the RCE state and compare it to the inital (standard) atmop
 
 ```{code-cell} ipython3
 fig, ax = plt.subplots()
-plots.profile_p_log(atmosphere['plev'], atmosphere['T'][-1, :])
+plots.profile_p_log(atmosphere['plev'], atmosphere['T'][-1, :], label="Init. state")
+plots.profile_p_log(rce.atmosphere['plev'], rce.atmosphere['T'][-1, :], label="RCE")
+ax.legend()
 ax.set_xlabel(r"$T$ / K")
 ax.set_ylabel("$p$ / hPa")
 ```

howto/humidity.md

@@ -58,13 +58,13 @@ $Q_r$ is a decisive quantity in climate science as it destabilizies the atmosphe
 
 ```{code-cell} ipython3
 fig, (ax0, ax1) = plt.subplots(ncols=2, sharey=True)
-plots.profile_p_log(atmosphere['plev'], atmosphere['H2O'][-1, :], ax=ax0)
+plots.profile_p_log(rce.atmosphere['plev'], rce.atmosphere['H2O'][-1, :], ax=ax0)
 ax0.set_xlabel(r"$q$ / VMR")
 ax0.set_xscale("log")
 ax0.set_ylabel("$p$ / hPa")
 
 ax1.axvline(0, color="k", linewidth=0.8)
-plots.profile_p_log(atmosphere['plev'], rce.radiation["net_htngrt"][-1, :], ax=ax1)
+plots.profile_p_log(rce.atmosphere['plev'], rce.radiation["net_htngrt"][-1, :], ax=ax1)
 ax1.set_xlim(-2, 0.5)
 ax1.set_xlabel(r"$Q_\mathrm{r}$ / (K/day)")
 ```
@@ -90,12 +90,12 @@ rce = konrad.RCE(
 rce.run()  # Start the simulation.
 
 fig, (ax0, ax1) = plt.subplots(ncols=2, sharey=True)
-plots.profile_p_log(atmosphere['plev'], atmosphere['H2O'][-1, :], ax=ax0)
+plots.profile_p_log(rce.atmosphere['plev'], rce.atmosphere['H2O'][-1, :], ax=ax0)
 ax0.set_xlabel(r"$q$ / VMR")
 ax0.set_ylabel("$p$ / hPa")
 
 ax1.axvline(0, color="k", linewidth=0.8)
-plots.profile_p_log(atmosphere['plev'], rce.radiation["net_htngrt"][-1, :], ax=ax1)
+plots.profile_p_log(rce.atmosphere['plev'], rce.radiation["net_htngrt"][-1, :], ax=ax1)
 ax1.set_xlim(-2, 0.5)
 ax1.set_xlabel(r"$Q_\mathrm{r}$ / (K/day)")
 ```

howto/lapserate.md

@@ -57,7 +57,7 @@ for lapserate in [6.5, 8, 10]:
     )
     rce.run()  # Start the simulation.
 
-    plots.profile_p_log(atmosphere['plev'], atmosphere['T'][-1, :])
+    plots.profile_p_log(rce.atmosphere['plev'], rce.atmosphere['T'][-1, :])
 ax.set_xlabel(r"$T$ / K")
 ax.set_ylabel("$p$ / hPa")
 ```
@@ -79,7 +79,7 @@ rce = konrad.RCE(
 rce.run()  # Start the simulation.
 
 fig, ax = plt.subplots()
-plots.profile_p_log(atmosphere['plev'], atmosphere['T'][-1, :])
+plots.profile_p_log(rce.atmosphere['plev'], rce.atmosphere['T'][-1, :])
 ax.set_xlabel(r"$T$ / K")
 ax.set_ylabel("$p$ / hPa")
 ```
@@ -106,7 +106,7 @@ for Ts in [280, 290, 300]:
     )
     rce.run()  # Start the simulation.
 
-    plots.profile_p_log(atmosphere['plev'], atmosphere['T'][-1, :])
+    plots.profile_p_log(rce.atmosphere['plev'], rce.atmosphere['T'][-1, :])
 ax.set_xlabel(r"$T$ / K")
 ax.set_ylabel("$p$ / hPa")
 ```

howto/netcdf_output.md

@@ -92,5 +92,5 @@ rce = konrad.RCE(
 rce.run()
 
 # Plot adjusted state
-plots.profile_p_log(atmosphere["plev"], atmosphere["T"][-1])
+plots.profile_p_log(rce.atmosphere["plev"], rce.atmosphere["T"][-1])
 ```

howto/ozone.md

@@ -44,15 +44,15 @@ for Ts in [285, 295, 305]:
     )
     rce.run()
 
-    l, = plots.profile_p_log(atmosphere["plev"], atmosphere["T"][-1], ax=ax0)
+    l, = plots.profile_p_log(rce.atmosphere["plev"], rce.atmosphere["T"][-1], ax=ax0)
     ax0.set_xlabel(r"$T$ / K")
     ax0.set_xlim(180, 306)
     ax0.set_ylabel("$p$ / hPa")
     ax0.set_ylim(bottom=atmosphere["plev"].max())
 
     plots.profile_p_log(
-        atmosphere["plev"],
-        atmosphere["O3"][-1] * 1e6,
+        rce.atmosphere["plev"],
+        rce.atmosphere["O3"][-1] * 1e6,
         label=f"{Ts} K",
         color=l.get_color(),
         ax=ax1,
@@ -81,15 +81,15 @@ for Ts in [285, 295, 305]:
     )
     rce.run()
 
-    l, = plots.profile_p_log(atmosphere["plev"], atmosphere["T"][-1], ax=ax0)
+    l, = plots.profile_p_log(rce.atmosphere["plev"], rce.atmosphere["T"][-1], ax=ax0)
     ax0.set_xlabel(r"$T$ / K")
     ax0.set_xlim(180, 306)
     ax0.set_ylabel("$p$ / hPa")
-    ax0.set_ylim(bottom=atmosphere["plev"].max())
+    ax0.set_ylim(bottom=rce.atmosphere["plev"].max())
 
     plots.profile_p_log(
-        atmosphere["plev"],
-        atmosphere["O3"][-1] * 1e6,
+        rce.atmosphere["plev"],
+        rce.atmosphere["O3"][-1] * 1e6,
         label=f"{Ts} K",
         color=l.get_color(),
         ax=ax1,

konrad/atmosphere.py

@@ -284,24 +284,6 @@ def refine_plev(self, phlev, **kwargs):
 
         return new_atmosphere
 
-    def copy(self):
-        """Create a copy of the atmosphere.
-
-        Returns:
-            konrad.atmosphere: copy of the atmosphere
-        """
-        datadict = dict()
-        datadict["phlev"] = copy(self["phlev"])  # Copy pressure grid.
-
-        # Create copies (and not references) of all atmospheric variables.
-        for variable in self.atmosphere_variables:
-            datadict[variable] = copy(self[variable]).ravel()
-
-        # Create a new atmosphere object from the filled data directory.
-        new_atmosphere = type(self).from_dict(datadict)
-
-        return new_atmosphere
-
     def calculate_height(self):
         """Calculate the geopotential height."""
         g = constants.earth_standard_gravity

konrad/component.py

@@ -1,3 +1,4 @@
+import copy
 import operator
 from collections.abc import Hashable
 
@@ -224,3 +225,7 @@ def get(self, variable, default=None, keepdims=True):
                 raise KeyError(f"'{variable}' not found and no default given.")
 
         return values if keepdims else values.ravel()
+
+    def copy(self):
+        """Return a deepcopy of the component."""
+        return copy.deepcopy(self)

konrad/core.py

@@ -154,7 +154,7 @@ def __init__(
         # Sub-model initialisation
 
         # Atmosphere
-        self.atmosphere = atmosphere
+        self.atmosphere = atmosphere.copy()
 
         # Radiation
         if radiation is None:
@@ -171,7 +171,7 @@ def __init__(
         # Surface
         self.surface = utils.return_if_type(
             surface, "surface", Surface, FixedTemperature()
-        )
+        ).copy()
 
         # Cloud
         self.cloud = utils.return_if_type(

konrad/test/test_atmosphere.py

@@ -57,3 +57,14 @@ def test_refine_plev(self, atmosphere_obj):
         atmosphere_new = atmosphere_obj.refine_plev(phlev=phlev)
 
         # TODO Add a proper test of values
+
+    def test_copy(self, atmosphere_obj):
+        """Test deepcopy of atmosphere component."""
+        atmosphere_copy =  atmosphere_obj.copy()
+
+        # Check if copied data is equal
+        assert np.array_equal(atmosphere_obj["T"], atmosphere_copy["T"])
+
+        # Check if copied data is independent
+        atmosphere_obj["T"][:] += 1.0
+        assert not np.array_equal(atmosphere_obj["T"], atmosphere_copy["T"])

konrad/test/test_surface.py

@@ -20,3 +20,15 @@ def test_init(self):
         for t in init_args:
             surf = surface.FixedTemperature(temperature=t)
             assert np.array_equal(surf["temperature"], np.array([280.0], dtype=float))
+
+    def test_copy(self):
+        """Test deepcopy of surface component."""
+        surf = surface.FixedTemperature(temperature=288.)
+        surf_copy =  surf.copy()
+
+        # Check if copied data is equal
+        assert np.array_equal(surf["temperature"], surf_copy["temperature"])
+
+        # Check if copied data is independent
+        surf["temperature"][:] += 1.0
+        assert not np.array_equal(surf["temperature"], surf_copy["temperature"])