`MeridionalHeatDiffusion` issue

[HenryDane]

Trying to create an EBM like:

import climlab
import numpy as np
import matplotlib.pyplot as plt

model_dims = {'lat':20, 'lon':10, 'lev':30}

state = climlab.volume_state(num_lat=model_dims['lat'], num_lon=model_dims['lon'], num_lev=model_dims['lev'])

ebm = climlab.EBM(state=state)

ebm.integrate_years(1)

leads to a crash here:

...
1 ebm = climlab.EBM(state=state)

File ~/Documents/_Feldl/climlab/climlab/model/ebm.py:271, in EBM.__init__(self, num_lat, num_lon, S0, s2, A, B, D, water_depth, Tf, a0, a2, ai, timestep, T0, T2, **kwargs)
    266 alb = albedo.StepFunctionAlbedo(state=self.state, **self.param)
    267 sw = SimpleAbsorbedShortwave(state=self.state,
    268                              insolation=ins.insolation,
    269                              albedo=alb.albedo,
    270                              **self.param)
--> 271 diff = MeridionalHeatDiffusion(state=self.state, use_banded_solver=False, **self.param)
    272 self.add_subprocess('LW', lw)
    273 self.add_subprocess('insolation', ins)

File ~/Documents/_Feldl/climlab/climlab/dynamics/meridional_heat_diffusion.py:74, in MeridionalHeatDiffusion.__init__(self, D, use_banded_solver, **kwargs)
     71 super(MeridionalHeatDiffusion, self).__init__(K=1.,
     72                 use_banded_solver=use_banded_solver, **kwargs)
     73 #  Now initialize properly
---> 74 self.D = D
     75 self.add_diagnostic('heat_transport', 0.*self.diffusive_flux)
     76 self.add_diagnostic('heat_transport_convergence', 0.*self.flux_convergence)

File ~/Documents/_Feldl/climlab/climlab/dynamics/meridional_heat_diffusion.py:84, in MeridionalHeatDiffusion.D(self, Dvalue)
     81 @D.setter
     82 def D(self, Dvalue):
     83     self._D = Dvalue
---> 84     self._update_diffusivity()

File ~/Documents/_Feldl/climlab/climlab/dynamics/meridional_heat_diffusion.py:90, in MeridionalHeatDiffusion._update_diffusivity(self)
     88     heat_capacity = value.domain.heat_capacity
     89 # diffusivity in units of m**2/s
---> 90 self.K = self.D / heat_capacity * const.a**2

File ~/Documents/_Feldl/climlab/climlab/dynamics/advection_diffusion.py:155, in AdvectionDiffusion.K(self, Kvalue)
    152 @K.setter  # currently this assumes that Kvalue is scalar or has the right dimensions...
    153 def K(self, Kvalue):
    154     self._K = Kvalue
--> 155     self._compute_advdiff_matrix()

File ~/Documents/_Feldl/climlab/climlab/dynamics/advection_diffusion.py:180, in AdvectionDiffusion._compute_advdiff_matrix(self)
    179 def _compute_advdiff_matrix(self):
--> 180     Karray = np.ones_like(self._Xbounds) * self.K
    181     try:
    182         Uarray = np.ones_like(self._Xbounds) * self.U

ValueError: operands could not be broadcast together with shapes (10,30,21) (30,) 

which is traced to self.K being set like:

    for varname, value in self.state.items():
        heat_capacity = value.domain.heat_capacity
    # diffusivity in units of m**2/s
    self.K = self.D / heat_capacity * const.a**2        

since heat_capacity is always created in the shape of elevation.

However, it seems that self.K should always have the shape of the acis which matches self.diffusion_axis (and also self.diffusion_axis_index).

[HenryDane]

On advection_diffusion.py:L180, it turns out that self._Xbounds has dimension of (lon, lev, lat+1) instead of (lat, lon, lev). The cause of this reordering is currently a mystery to me.

[HenryDane]

I don't know what was happening here -- it got sorted out on the diff_lon branch.