-
Notifications
You must be signed in to change notification settings - Fork 108
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
Caching parsed output files #273
Changes from all commits
7fbf27e
856e650
58aca7f
9335f94
16d745e
de7bd47
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,38 @@ | ||
import unittest | ||
|
||
from custodian.utils import tracked_lru_cache | ||
|
||
|
||
class TrackedLruCacheTest(unittest.TestCase): | ||
def setUp(self): | ||
# clear cache before and after each test to avoid | ||
# unexpected caching from other tests | ||
tracked_lru_cache.tracked_cache_clear() | ||
|
||
def test_cache_and_clear(self): | ||
n_calls = 0 | ||
|
||
@tracked_lru_cache | ||
def some_func(x): | ||
nonlocal n_calls | ||
n_calls += 1 | ||
return x | ||
|
||
assert some_func(1) == 1 | ||
assert n_calls == 1 | ||
assert some_func(2) == 2 | ||
assert n_calls == 2 | ||
assert some_func(1) == 1 | ||
assert n_calls == 2 | ||
|
||
assert len(tracked_lru_cache.cached_functions) == 1 | ||
|
||
tracked_lru_cache.tracked_cache_clear() | ||
|
||
assert len(tracked_lru_cache.cached_functions) == 0 | ||
|
||
assert some_func(1) == 1 | ||
assert n_calls == 3 | ||
|
||
def tearDown(self): | ||
tracked_lru_cache.tracked_cache_clear() |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,38 @@ | ||
""" | ||
Helper functions for dealing with vasp files. | ||
""" | ||
|
||
from pymatgen.io.vasp.outputs import Outcar, Vasprun | ||
|
||
from custodian.utils import tracked_lru_cache | ||
|
||
|
||
@tracked_lru_cache | ||
def load_vasprun(filepath, **vasprun_kwargs): | ||
""" | ||
Load Vasprun object from file path. | ||
Caches the output for reuse. | ||
|
||
Args: | ||
filepath: path to the vasprun.xml file. | ||
**vasprun_kwargs: kwargs arguments passed to the Vasprun init. | ||
|
||
Returns: | ||
The Vasprun object | ||
""" | ||
return Vasprun(filepath, **vasprun_kwargs) | ||
|
||
|
||
@tracked_lru_cache | ||
def load_outcar(filepath): | ||
""" | ||
Load Outcar object from file path. | ||
Caches the output for reuse. | ||
|
||
Args: | ||
filepath: path to the OUTCAR file. | ||
|
||
Returns: | ||
The Vasprun object | ||
""" | ||
return Outcar(filepath) |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -12,3 +12,13 @@ def _patch_get_potential_energy(monkeypatch): | |
Monkeypatch the multiprocessing.cpu_count() function to always return 64 | ||
""" | ||
monkeypatch.setattr(multiprocessing, "cpu_count", lambda *args, **kwargs: 64) | ||
|
||
|
||
@pytest.fixture(autouse=True) | ||
def _clear_tracked_cache(): | ||
""" | ||
Clear the cache of the stored functions between the tests. | ||
""" | ||
from custodian.utils import tracked_lru_cache | ||
|
||
tracked_lru_cache.tracked_cache_clear() | ||
Comment on lines
+17
to
+24
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think this fixture currently runs on every test, not just the handler tests which would be unnecessary work? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think you can move it to There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Since this is also used in the validators, I think the fixture should either be duplicated in the handlers and validators files, or added it explicitly to the tests that need it, removing the There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yeah, probably duplicating it in both files makes it more explicit too. People are more likely to see it and become aware this is happening. |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,27 @@ | ||
import os | ||
import unittest | ||
|
||
from custodian.utils import tracked_lru_cache | ||
from custodian.vasp.io import load_outcar, load_vasprun | ||
|
||
test_dir = os.path.join(os.path.dirname(__file__), "..", "..", "..", "test_files") | ||
|
||
|
||
class IOTest(unittest.TestCase): | ||
def test_load_outcar(self): | ||
outcar = load_outcar(os.path.join(test_dir, "large_sigma", "OUTCAR")) | ||
assert outcar is not None | ||
outcar2 = load_outcar(os.path.join(test_dir, "large_sigma", "OUTCAR")) | ||
|
||
assert outcar is outcar2 | ||
|
||
assert len(tracked_lru_cache.cached_functions) == 1 | ||
|
||
def test_load_vasprun(self): | ||
vr = load_vasprun(os.path.join(test_dir, "large_sigma", "vasprun.xml")) | ||
assert vr is not None | ||
vr2 = load_vasprun(os.path.join(test_dir, "large_sigma", "vasprun.xml")) | ||
|
||
assert vr is vr2 | ||
|
||
assert len(tracked_lru_cache.cached_functions) == 1 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
How about instead of a new
io
module, we just definedirectly in
custodian/vasp/handlers.py
?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This could be done, but since this is also used in the
validators
module they should be imported from thehandlers
. I think that having a generic io module could be fine, but I have no problem moving them to handlers.py.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, I guess no need to move to
handlers.py
.