[Feature Request] Give users finer control over chunking when computing optical properties

[dhjx1996]

Is your feature request related to a problem? Please describe.
When calling

clouds_optical_props = cloud_optics.compute_cloud_optics(
    atmosphere, problem_type="two-stream",
)

the chunking of the resultant clouds_optical_props is different from the chunking of atmosphere and that seems to be the cause of memory problems downstream.

Describe the solution you'd like
The chunking of atmosphere should be maintained and carried over into clouds_optical_props (I suggest the same for compute_gas_optics though I did not test that function).

Describe alternatives you've considered
Alternatively, give users finer control over the chunking in compute_cloud_optics, compute_gas_optics and similar functions.

Additional context
Screenshots of the situation I encountered are attached. Notice that (Ydim, Xdim) of atmosphere is (180, 180) whereas (Ydim, Xdim) of clouds_optical_props is (12, 2880).

[brendancol]

@dhjx1996 thanks so much for testing and reporting.

Are you installing from the main branch, or using the 0.0.9 release? We need to get a more recent release as there has been some changes...nothing that would effect what you are reporting though.

While users can chunk / rechunk as wanted, I think it makes sense to use the chunking from atmosphere.

Is it fair to say this wouldn't be "finer grain control", but instead enforcing the chunking from atmosphere?

[dhjx1996]

Hi @brendancol, thanks for the quick reply! I am using conda to install and the version is 0.0.9. I will say that my concern is less with the exact chunking -- since I can easily rechunk the results -- but more of the overall memory management in compute_cloud_optics (and compute_gas_optics). I have been facing memory spikes upwards of 128gb (which crashes the server I am using) when I run dask.compute or dask.persist downstream of compute_cloud_optics. I am still performing some tests (it is difficult to pinpoint errors in Dask due to lazy evaluations) but I intend to create a separate issue on that soon.

[sehnem]

@dhjx1996 there was a fix in dask 2025.3 that may solve the memory issues, I created a PR for it #201

[brendancol]

@dhjx1996 There is a PR here with a dask.ipynb notebook which I ran on an AWS r5a.16xlarge instance

[dhjx1996]

@dhjx1996 There is a PR here with a dask.ipynb notebook which I ran on an AWS r5a.16xlarge instance

Thanks! I'll take a look

[dhjx1996]

@dhjx1996 There is a PR here with a dask.ipynb notebook which I ran on an AWS r5a.16xlarge instance

@RobertPincus