Enable lazy loading of submodules

[arjxn-py]

Enabling lazy loading of submodules should potentially make PyBaMM import faster. See also: https://scientific-python.org/specs/spec-0001/

Originally posted by @agriyakhetarpal in https://github.com/pybamm-team/PyBaMM/issues/3475#issuecomment-1786894302

[prady0t]

Hey can I look into it?

[agriyakhetarpal]

Hi @prady0t, thanks for helping us. Doing this would be much appreciated, but it would be better to start after #3475 gets completed. While we need import pybamm to be faster, the expected outcome is that there must not be any changes to the user-facing API.

[prady0t]

Thanks for replying. I will be looking at other issues in the meantime.

[agriyakhetarpal]

See also: https://docs.python.org/3/library/importlib.html#implementing-lazy-imports; a rather crude, bare-bones implementation for delaying imports

[arjxn-py]

I can look into this now. Keeping in mind that it is not a beginner level issue and @prady0t is occupied with other one.

[AbhishekChaudharii]

Hi @arjxn-py @agriyakhetarpal If you're not working on this issue, May I go ahead and try to solve this? Also can you guide me?

[arjxn-py]

Hey @AbhishekChaudharii, thanks for showing your interest on this but I am currently on it and going to open up a PR in a couple of days. Feel free to look into other issues 🙂

[AbhishekChaudharii]

Sure no problem :)

[kratman]

@arjxn-py Can we close this now?

[agriyakhetarpal]

I would say we should keep it open (other, large packages are able to make it work – maybe we can too), but yes, it is @arjxn-py's call after all

[kratman]

I thought the consensus was that it was not a major bottleneck so the work/code required to do it was bigger than the problem

[agriyakhetarpal]

Right, maybe let's not close it because it's still within scope but lower down the priority for now – I'll change the labels.

[kratman]

We can always open new tickets when problems recur, no need to keep this open

[agriyakhetarpal]

But this is a recurring problem, isn't it? I just tested this in a fresh notebook/kernel without sys.modules being cached and it took ~20 seconds to import PyBaMM – so it's probably still valid. At least for organisational purposes it would be nice to keep open

[agriyakhetarpal]

I just noticed that https://github.com/scientific-python/lazy_loader is now being recommended for Python versions 3.11 and later, and that too, with very specific minimum patch versions to avoid races, say 3.11.9, for example. Python 3.11 being the lowest version for us will be quite many years away at the time of writing (at least four). It sucks that I would have liked to have seen this done for PyBaMM but I don't think we have a viable choice at this time, I think we can close this and re-open later (or open a new issue, whichever way works). A better and more reliable way would be to profile PyBaMM's import and see which of the bigger modules are causing the bottleneck, and if needed then implement minor breaking changes or lazy loading at the functionality level manually at the time of use of a function, so that we aren't importing a lot of those functions when one runs import pybamm. So I'm agreeing with @kratman here and closing this issue, please feel more than free to re-open later, @arjxn-py, or to suggest an alternative enhancement to speed up the import (whose slowness remains very much a valid problem – it's just not suited to be done for the entire public API because of problems that have been previously discussed).