Skip to content

Commit

Permalink
more lint
Browse files Browse the repository at this point in the history
  • Loading branch information
gpetretto committed Apr 17, 2023
1 parent ebe4411 commit 676432a
Showing 1 changed file with 45 additions and 16 deletions.
61 changes: 45 additions & 16 deletions src/atomate2/abinit/powerups.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,35 @@
from typing import Any

from jobflow import Flow, Job, Maker
from pymatgen.io.abinit import PseudoTable
from pymatgen.io.abinit.abiobjects import KSampling

from atomate2.abinit.jobs.base import BaseAbinitMaker


def update_maker_kwargs(class_filter, dict_mod_updates, flow, name_filter):
"""
Update an object inside a Job, a Flow or a Maker.
A generic method to be shared for more specific updates that will
build the dict_mod_updates.
Parameters
----------
flow : .Job or .Flow or .Maker
A job, flow or Maker.
dict_mod_updates : dict
The updates to apply.
name_filter : str or None
A filter for the name of the jobs.
class_filter : Maker or None
A filter for the class used to generate the flows. Note the class
filter will match any subclasses.
Returns
-------
Job or Flow or Maker
A copy of the input modified flow/job/maker.
"""
updated_flow = deepcopy(flow)
if isinstance(updated_flow, Maker):
updated_flow = updated_flow.update_kwargs(
Expand Down Expand Up @@ -56,8 +78,8 @@ def update_user_abinit_settings(
name_filter : str or None
A filter for the name of the jobs.
class_filter : Maker or None
A filter for the BaseAbinitMaker class used to generate the flows. Note the class
filter will match any subclasses.
A filter for the BaseAbinitMaker class used to generate the flows. Note the
class filter will match any subclasses.
Returns
-------
Expand All @@ -68,7 +90,9 @@ def update_user_abinit_settings(
f"input_set_generator->user_abinit_settings->{k}": v
for k, v in abinit_updates.items()
}
updated_flow = update_maker_kwargs(class_filter, dict_mod_updates, flow, name_filter)
updated_flow = update_maker_kwargs(
class_filter, dict_mod_updates, flow, name_filter
)
return updated_flow


Expand Down Expand Up @@ -97,8 +121,8 @@ def update_factory_kwargs(
name_filter : str or None
A filter for the name of the jobs.
class_filter : Maker or None
A filter for the BaseAbinitMaker class used to generate the flows. Note the class
filter will match any subclasses.
A filter for the BaseAbinitMaker class used to generate the flows. Note the
class filter will match any subclasses.
Returns
-------
Expand All @@ -109,7 +133,9 @@ def update_factory_kwargs(
f"input_set_generator->factory_kwargs->{k}": v
for k, v in factory_updates.items()
}
updated_flow = update_maker_kwargs(class_filter, dict_mod_updates, flow, name_filter)
updated_flow = update_maker_kwargs(
class_filter, dict_mod_updates, flow, name_filter
)
return updated_flow


Expand Down Expand Up @@ -139,8 +165,8 @@ def update_user_kpoints_settings(
name_filter : str or None
A filter for the name of the jobs.
class_filter : Maker or None
A filter for the BaseAbinitMaker class used to generate the flows. Note the class
filter will match any subclasses.
A filter for the BaseAbinitMaker class used to generate the flows. Note the
class filter will match any subclasses.
Returns
-------
Expand All @@ -157,7 +183,9 @@ def update_user_kpoints_settings(
for k, v in kpoints_updates.items()
}

updated_flow = update_maker_kwargs(class_filter, dict_mod_updates, flow, name_filter)
updated_flow = update_maker_kwargs(
class_filter, dict_mod_updates, flow, name_filter
)
return updated_flow


Expand All @@ -181,21 +209,22 @@ def update_generator_attributes(
flow : .Job or .Flow or .Maker
A job, flow or Maker.
generator_updates : dict
The updates to apply to the input genrator.
The updates to apply to the input generator.
name_filter : str or None
A filter for the name of the jobs.
class_filter : Maker or None
A filter for the BaseAbinitMaker class used to generate the flows. Note the class
filter will match any subclasses.
A filter for the BaseAbinitMaker class used to generate the flows. Note the
class filter will match any subclasses.
Returns
-------
Job or Flow or Maker
A copy of the input flow/job/maker modified to use the updated factory settings.
"""
dict_mod_updates = {
f"input_set_generator->{k}": v
for k, v in generator_updates.items()
f"input_set_generator->{k}": v for k, v in generator_updates.items()
}
updated_flow = update_maker_kwargs(class_filter, dict_mod_updates, flow, name_filter)
updated_flow = update_maker_kwargs(
class_filter, dict_mod_updates, flow, name_filter
)
return updated_flow

0 comments on commit 676432a

Please sign in to comment.