Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

BUMBANUMBA #36

Merged
merged 8 commits into from
Jun 3, 2020
Merged

BUMBANUMBA #36

merged 8 commits into from
Jun 3, 2020

Conversation

ericbarefoot
Copy link
Collaborator

@ericbarefoot ericbarefoot commented May 22, 2020

Numba speedup

I was interested in what @amoodie did in #27, so I started messing around with it too once #29 was merged. In particular, with more modular functions, how could a lot of things be moved into a utils-like module, jitted, and then called that way.

I basically took a profile of the code, and just started jitting the things that took the longest, to see how fast we could get. So far, I have pulled out and jitted the following:

  • the construction of the water weight array
  • calculating the next step indices for water
  • updating the water flow field.
  • checking for loops in the water paths
  • the random pick function
  • choosing weights for sediment routing
  • spreading sand between cells.

There are quite a few more things that take a significant fraction of the computational time, and could also go into jitted functions. However, I'm stopping here for now, because with the default parameters, I see about a 25x speedup from the develop branch with these adjustments. Using this script, I get about 9.54s per iteration for develop, and 0.36s per iteration after jitted functions are compiled.

from pyDeltaRCM.deltaRCM_driver import pyDeltaRCM
import time

delta = pyDeltaRCM()

start1 = time.time()
delta.update()
end1 = time.time()
print("Elapsed time for first iteration including compile = %s" % str((end1 - start1)))

n = 10
start = time.time()
for _t in range(0, n):
    delta.update()
    print("timestep:", _t)

end = time.time()

print("Elapsed time per iteration with jitting = %s" % str((end - start) / n))

Needs

I'd like y'all's feedback on the organization. Basically, I'm not sure about placing these things in a separate utils, so maybe you can comment @amoodie, if it would be good to move the things that belong to water into the water_tools file, but have them as standalone jitted functions in the water_tools module.

Some other ideas I have that may make this organization more intuitive and the long function calls more palatable:

  • rename shared_tools to utils like @amoodie did
  • move some commonly used attributes of the classes that are not set by the user to the utils module as getter functions. I kinda did this for the self.iwalk and self.jwalk attributes, which are initialized now using a getter function from shared_tools. The idea is that if these getters become compiled themselves, we could have them called inside jitted functions and not have to supply them in the arguments.
  • move jitted functions unique to sed and water routing into the module where these classes are defined.

Thoughts??

@ericbarefoot ericbarefoot added the enhancement New feature or request label May 22, 2020
@ericbarefoot ericbarefoot changed the title Numbajit BUMBANUMBA May 22, 2020
@ericbarefoot
Copy link
Collaborator Author

Ok. I also need some advice with the docs, because that's what failed the second time around. How should I fix that?

@ericbarefoot ericbarefoot linked an issue May 22, 2020 that may be closed by this pull request
@ericbarefoot ericbarefoot marked this pull request as draft May 22, 2020 19:08
@ericbarefoot
Copy link
Collaborator Author

I marked this as a draft because of two things. I will have to learn a little bit more about the summary for modules so that the module with the jitted functions becomes documented. Also, it looks like this branch is not producing the correct results. I just looked at a picture of the results two steps in, and something is different than develop. However, I still want feedback on other stuff while I hunt down these bugs.

@amoodie
Copy link
Member

amoodie commented May 23, 2020

Hi @ericbarefoot, awesome effort.

Also, it looks like this branch is not producing the correct results. I just looked at a picture of the results two steps in, and something is different than develop.

Yeah, I can verify this branch produces different results, based on my PR for a consistency test. I thought maybe this was because your branch uses a non-jitted pick_random_inlet function, which would change the number of random calls, but this does not appear to be the case (see below).

I don't think it matters what you decide to name the file with all these new jitted functions. It could make sense to change the name since shared_tools is no longer inherited by Tools like all of the other _tools methods. I think shared would a fine and general-enough module name.

I'm gonna send you a PR now that jits the pick_random_inlet and fixes the documentation with all the new jitted functions you have added. If you change the name of shared_tools, just update the doc folder structure and files accordingly. You can always test the docs locally with (cd docs && make test) from the root folder, or test and build the docs with (cd docs && make docs) (see the targets in the Makefile).

@amoodie
Copy link
Member

amoodie commented May 24, 2020

hey @ericbarefoot, one issue might be that the changes you made to random_pick now allow the choice to the same cell (i.e., the center of the 3x3 neighborhood). I think before, when the method took an 8 cell array, the choice returned was never the same cell as current.

Maybe you can try to revert this behavior and see if you get a consistent result with develop? You'll need to use my jitted random_pick_inlet too, for sure.

Nevermind, I think that part actually is fine lol

@ericbarefoot
Copy link
Collaborator Author

ericbarefoot commented May 26, 2020

Hey @amoodie I just went spelunking to find the bugs that were messing with water and sediment routing. I found and fixed them, but it affects some of the things in your PR, which I didn't want to change further to find the problem.

Can you rebase your branch on ef18f96? This is a rebase of the changes I made on the new stuff in develop. That way I can easily include your fixes to the docs and the other list typing stuff.

This makes them available outside the object, which is important in some
cases, and the class just gets them from here.

feat: make shared_tools into a module, not a class.

This is in preparation for jitting, since we want them to be standalone
functions.

feat: break out step that updates water parcel positions

These functions now become helper utility functions in shared_tools, and
are wrapped by update_Q, which handles them all, and passes them what
they need.

refactor: break out sed functions into smaller functions

This modularizes the code and makes it easier to compile.

feat: jit the functions

apply jit to functions, and fix some type problems that came along with
refactoring.

refactor: move next step calculation to jitted functions

The next slowest operation was getting the indexes of the next cells, so
I moved that to a jitted function too.

refactor: jit the calculation of weight array for the water

fix: move random number generation into numba

Copy of what @amoodie did in DeltaRCM#27

fix: some broken tests to test the shared tools portion.

fix: add numba to requirements

fix: move all inlet picking into the random_pick function

There's essentially no difference between the random_pick function and
the random_pick_inlet function other than there is no preference given
for the inlet cells, so you can pick whichever one, so I just adjust the
mechanism to use the existing random pick machinery.

fix: water and sediment routing bugs that caused inconsistent behavior

There were two issues with the refactoring. One was that water routing
developed a flat flow field because I wasn't adding each passing water
parcel to the Q field. The other was related to sediment parcel routing,
where I had transcribed the i and j of stepping incorrectly in the
jitted functions.
Andrew Moodie and others added 7 commits May 26, 2020 13:50
@ericbarefoot ericbarefoot marked this pull request as ready for review June 2, 2020 23:01
@ericbarefoot
Copy link
Collaborator Author

Hey all: I think this is ready to go now. I wrote a whole bunch of tests for the new jitted functions today. Turns out, the coveralls doesn't say lines are "covered" if the code is compiled prior to running, so we can think about how to address that if you all would like to make sure we continue to increase coverage always.

@ericbarefoot ericbarefoot mentioned this pull request Jun 3, 2020
5 tasks
Copy link
Member

@amoodie amoodie left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried messing around with an environment variable to the coverage job on travis (NUMBA_DISABLE_JIT) I thought would fix this coverage issue, but apparently it does not, based on some comments around github. I think we could open an issue to keep checking on it and see if a fix comes along.

Based on 092d98c, you were able to fix the bug without changing the test_consistent? Nice. And did you run that yaml configuration I had given you to verify the bug is resolved?

Otherwise, looks good to me 👍

@ericbarefoot
Copy link
Collaborator Author

Based on 092d98c, you were able to fix the bug without changing the test_consistent?

As far as I can tell, it's close enough yeah. I was surprised too.

And did you run that yaml configuration I had given you to verify the bug is resolved?

Yup. no more strange zig zagging paths, and it makes it through without erroring out.

@amoodie
Copy link
Member

amoodie commented Jun 3, 2020

As far as I can tell, it's close enough yeah. I was surprised too.

Thinking about it a bit more, I guess it makes sense. The test never had the bug (only ran for one small timestep), so why would it be affected by the bug being fixed‽

Yup. no more strange zig zagging paths, and it makes it through without erroring out.

Cool, anyway yeah lgtm, not sure if Jay wants to look or not.

Copy link
Member

@elbeejay elbeejay left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great effort! This is all really clean and easy to follow.

More a note than any commentary on this PR, but in general we should probably try running longer "test" simulations for large changes to check for longer-term stability as those types of checks are not going to be possible via unit testing or any continuous integration method.

self.iwalk = np.array([[-1, 0, 1],
[-1, 0, 1],
[-1, 0, 1]])
self.iwalk = shared_tools.get_iwalk()

self.jvec = np.array([[-sqrt05, -1, -sqrt05],
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If iwalk and jwalk are moved to shared_tools (which is fine by me) does it make sense to also move over all of the other walk/vector arrays like jvec and so forth?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I think I see a larger 'reorganize and tidy' PR on the horizon, where we once again move stuff into separate classes. I think this relates to the milestones for splitting out the hydrodynamics and morphodynamics code into separate modular classes, with a utils or similar module that they both access. I think moving a bunch of the stuff like the definitions of these arrays into something like utils makes a lot of sense.

@elbeejay
Copy link
Member

elbeejay commented Jun 3, 2020

Oh, and regarding the coverage not including compiled code I also did a cursory search and did not find a solution that looks like it will work for us. The closest thing I found was this, which pretty much required writing a numba toggle for each function and testing them with numba turned off to "cover" the code in the unit tests. For this project I don't know of a good reason why one would want to turn numba off so I think a toggle would just clutter things up. So I'm good on the coverage front as well.

@amoodie amoodie merged commit a7b6124 into DeltaRCM:develop Jun 3, 2020
@ericbarefoot ericbarefoot mentioned this pull request Jun 3, 2020
amoodie added a commit to amoodie/pyDeltaRCM that referenced this pull request Feb 25, 2021
@ericbarefoot ericbarefoot deleted the numbajit branch August 31, 2021 14:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Runs slowly
3 participants