Handle land mask

[paigem]

This PR does the following:

• obtains land mask info from sst field (currently assumes that the land is masked by NaN)
• shrinks the input arrays before sending to Fortran code
  • flattens the input arrays
  • removes land points from all variables (this should significantly help with speed up of runtime)
• unshrinks the output from the Fortran code and puts the NaNs back on land points

Additionally:

• removed the threadsafe for xarray, skin computations (see Attempt to fix threadsafe on skin #39 for more discussion on ways around this)
• added the bug fix in create_data() from Bugfix for test data #37
• new tests were added: test_flux_np.py
• separated out the creation of test data to its own file to avoid repetition: create_test_data.py

Closes #31

[jbusecke]

Nice work @paigem. This will be really useful! I had a few nitpicks and I think you should add some wording to the docstrings, but after that this seems good to go for me.

[jbusecke]

It might be helpful if @rabernat could look at the 'shrink'/'unshrink' logic and see if there are performance improvements that we have not considered?

Maybe we should profile the memory use of this?

[paigem]

I did some quick profiling to compare runtime comparing with and without the mask. This preliminary comparison shows a speed up of about 20% by masking land values.

This is not a complete comparison. The following assumptions are made:

• input size of (3600, 2700, 2) (CM2.6-sized with 2 time steps)
• masking 30% of input values
• using ecmwf algorithm (completely arbitrary)
  • I did a couple runs with coare3p0 and saw similar results (~20% speedup)
• only tried on noskin() so far

Snakeviz visualizations

No mask, CM2.6-sized test data:

data = create_data((3600, 2700, 2), chunks=None)
%snakeviz out_data = noskin_nomask(*data, 'ecmwf', 2, 10, 6)

We can see that the numpy wrapper noskin_np() takes essentially the entire runtime.

With mask, CM2.6-sized test data

data = create_data((3600, 2700, 2), chunks=None, land_mask=True)
%snakeviz out_data = noskin(*data, 'ecmwf', 2, 10, 6)

The small rectangles in the bottom right are the extra work needed to shrink and unshrink the arrays. Even with these extra computations, the runtime is noticeably shorter.

Runtime comparison

# Take average runtime across 3 runs

# No mask
nomask_avg = np.mean((60.467731952667236,61.46433997154236,61.295777797698975))

# Mask
mask_avg = np.mean((49.637826919555664,46.95607113838196,48.179039001464844))

total_percent_faster = (nomask_avg - mask_avg)/nomask_avg
total_percent_faster

Result: 0.2098748237283278 --> ~20% faster!

As @jbusecke stated above, we may also want to profile memory usage.

[jbusecke]

Thats amazing @paigem! Thanks for the nice comparison.

[jbusecke]

I had one more thing, that might be useful to change here. I believe now that we are converting the input to a 1d array in any case, we can get rid of this wrapper entirely.

To check that this is correct, I would write a test that puts arrays of various dimensions (1-4D?) through the xarray wrappers and make sure that this does not lead to crashes.

[paigem]

@jbusecke Take a look at the new test I wrote in this commit to verify that our xarray wrapper takes arrays of size 2d and greater (i.e. no longer just 3d!). It seems to work fine, but I wasn't sure how to write a test for that...

[jbusecke]

I am not sure I get this. If the fortran code crashes, this will never be called? I would remove it. We just need to make sure that we check the CI actions carefully.

[paigem]

Ok cool, I was under the impression we needed some sort of "check" (eg an assert statement). I can remove the unnecessary assert line.

[jbusecke]

Suggested change

	def test_all_input_array_sizes_valid(skin_correction):
	shapes = (
	(3, 4),
	(2, 3, 4),
	(2, 3, 4, 5),
	) # create_data() only allows for inputs of 2 or more dimensions
	data = (create_data(s, skin_correction=skin_correction) for s in shapes)
	if skin_correction:
	func = skin
	else:
	func = noskin
	tuple(func(*d, "coare3p0", 2, 10, 6) for d in data)
	assert (
	1 == 1
	) # This line is always true, but verifies that the above line doesn't crash the Fortran code
	@pytest.mark.parametrize('shape', [(3, 4), (2, 3, 4), (2, 3, 4, 5),])
	def test_all_input_array_sizes_valid(skin_correction, shape):
	# create_data() only allows for inputs of 2 or more dimensions
	data = create_data(shape, skin_correction=skin_correction)
	if skin_correction:
	func = skin
	else:
	func = noskin
	func(*data, "coare3p0", 2, 10, 6)

[jbusecke]

What I did here is factor out the shape as a parameterized input, so that each shape gets its own test. This enables a more fine grained control (e.g. if for some reason the 4d case fails, but the others pass we will see this immediately in the test report).

[jbusecke]

I made a few minor suggestions. If you agree with those, you can commit and merge once the tests pass (make sure to double check that the % line goes to 100 in the CI log).