Skip to content

Commit dbc626e

Browse files
mkphuthimkphuthi
committed
Added group_size tests
1 parent 62f06ec commit dbc626e

File tree

5 files changed

+125
-69
lines changed

5 files changed

+125
-69
lines changed

‎asimtools/asimmodules/active_learning/ase_md.py

+1
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,7 @@ def ase_md(
221221
friction=friction,
222222
)
223223
elif dynamics == 'npt':
224+
assert pfactor is not None, 'Pressure factor must be provided'
224225
atoms, _ = npt(
225226
atoms,
226227
temp,

‎asimtools/asimmodules/active_learning/compute_deviation.py

+28-7
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ def compute_deviation(
3636
:type calc_ids: Optional[Sequence]
3737
3838
"""
39-
properties = ['energy', 'forces', 'stress']
39+
properties = ['energy', 'forces', 'stress', 'energy_per_atom']
4040
if calc_ids is None:
4141
model_weights_files = natsorted(glob(model_weights_pattern))
4242

@@ -81,19 +81,25 @@ def compute_deviation(
8181

8282
atoms.set_calculator(calc)
8383
energy = atoms.get_potential_energy(atoms)
84+
energy_per_atom = energy / len(atoms)
8485
forces = np.linalg.norm(atoms.get_forces(), axis=1)
8586
stress = -np.sum(
8687
atoms.get_stress(voigt=True, include_ideal_gas=False)[:3]
8788
) / 3
8889

90+
prop_dict['energy_per_atom'][calc_id].append(energy_per_atom)
8991
prop_dict['energy'][calc_id].append(energy)
9092
prop_dict['forces'][calc_id].append(forces)
9193
prop_dict['stress'][calc_id].append(stress)
94+
atom_results['energy_per_atom'].append(energy_per_atom)
9295
atom_results['energy'].append(energy)
9396
atom_results['forces'].append(forces)
9497
atom_results['stress'].append(stress)
9598

9699
prop_dict['energy']['std'].append(np.std(atom_results['energy']))
100+
prop_dict['energy_per_atom']['std'].append(
101+
np.std(atom_results['energy_per_atom'])
102+
)
97103
prop_dict['forces']['mean_std'].append(
98104
np.mean(np.std(atom_results['forces'], axis=1))
99105
)
@@ -104,24 +110,39 @@ def compute_deviation(
104110

105111
df = pd.DataFrame({
106112
'energy_std': prop_dict['energy']['std'],
113+
'energy_per_atom_std': prop_dict['energy_per_atom']['std'],
107114
'force_mean_std': prop_dict['forces']['mean_std'],
108115
'force_max_std': prop_dict['forces']['max_std'],
109116
'stress_std': prop_dict['stress']['std']
110117

111118
})
112119
df.to_csv('deviations.csv')
113120

121+
unit_dict = {
122+
'energy': 'eV',
123+
'energy_per_atom': 'eV/atom',
124+
'forces': 'eV/$\AA$',
125+
'stress': 'eV/$\AA^3$',
126+
}
114127
for prop in properties:
115128
if prop not in ['forces']:
116129
# df = pd.DataFrame(prop_dict[prop])
117130
fig, ax = plt.subplots()
118131
for calc_id in calc_dict:
119132
ax.plot(prop_dict[prop][calc_id], label=calc_id)
120-
ax.set_xlabel('Image index')
121-
ax.set_ylabel(f'{prop} [ASE units]')
122-
ax.legend()
123-
plt.savefig(f'{prop}.png')
124-
plt.close()
133+
ax.set_ylabel(f'{prop} {unit_dict[prop]}')
134+
else:
135+
fig, ax = plt.subplots()
136+
for calc_id in calc_dict:
137+
ax.plot(
138+
np.max(prop_dict[prop][calc_id],axis=1),
139+
label=calc_id,
140+
)
141+
ax.set_ylabel(f'{prop} max {unit_dict[prop]}')
142+
ax.set_xlabel('Image index')
143+
ax.legend()
144+
plt.savefig(f'{prop}.png')
145+
plt.close()
125146

126147
fig, ax = plt.subplots()
127148
if prop == 'forces':
@@ -131,7 +152,7 @@ def compute_deviation(
131152
else:
132153
ax.plot(prop_dict[prop]['std'])
133154
ax.set_xlabel('Image index')
134-
ax.set_ylabel(f'{prop} std [ASE units]')
155+
ax.set_ylabel(f'{prop} std {unit_dict[prop]}')
135156
plt.savefig(f'{prop}_std.png')
136157
plt.close()
137158

‎asimtools/job.py

+37-31
Original file line numberDiff line numberDiff line change
@@ -632,28 +632,27 @@ def _gen_array_script(
632632
txt += '\necho "Job started on `hostname` at `date`"\n'
633633
txt += 'CUR_DIR=`pwd`\n'
634634
txt += 'echo "LAUNCHDIR: ${CUR_DIR}"\n'
635-
txt += f'GROUP_SIZE={group_size}\n'
636-
seqtxt = '$(seq $((${SLURM_ARRAY_TASK_ID}*${GROUP_SIZE})) '
637-
seqtxt += '$(((${SLURM_ARRAY_TASK_ID}+1)*${GROUP_SIZE})))'
635+
txt += f'G={group_size} #Group size\n'
636+
txt += 'N=${SLURM_ARRAY_TASK_ID}\n'
638637
txt += f'WORKDIRS=($(ls -dv ./id-*))\n'
638+
seqtxt = '$(seq $(($G*$N)) $(($G*$N+$G-1)) )'
639639
txt += f'for i in {seqtxt}; do\n'
640-
txt += 'cd ${WORKDIRS[$i]};\n'
640+
txt += ' WORKDIR=${WORKDIRS[$i]}\n'
641+
txt += ' cd ${WORKDIR};\n'
641642
# else:
642643
# txt += '\nif [[ ! -z ${SLURM_ARRAY_TASK_ID} ]]; then\n'
643644
# txt += ' fls=( id-* )\n'
644645
# txt += ' WORKDIR=${fls[${SLURM_ARRAY_TASK_ID}]}\n'
645646
# txt += 'fi\n\n'
646647
# txt += 'cd ${WORKDIR}\n'
647-
txt += '\n'
648-
txt += '\n'.join(slurm_params.get('precommands', []))
649-
txt += '\n'
650-
txt += '\n'.join(self.unitjobs[0].calc_params.get('precommands', []))
651-
txt += '\n'
652-
txt += 'echo "WORKDIR: ${WORKDIRS[$i]}"\n'
653-
txt += self.unitjobs[0].gen_run_command() + '\n'
654-
txt += '\n'.join(slurm_params.get('postcommands', []))
655-
txt += '\n'
656-
txt += 'cd ${CUR_DIR}\n'
648+
txt += ' ' + '\n'.join(slurm_params.get('precommands', []))
649+
txt += ' ' + '\n'.join(
650+
self.unitjobs[0].calc_params.get('precommands', [])
651+
)
652+
txt += ' echo "WORKDIR: ${WORKDIR}"\n'
653+
txt += ' ' + self.unitjobs[0].gen_run_command() + '\n'
654+
txt += ' ' + '\n'.join(slurm_params.get('postcommands', [])) + '\n'
655+
txt += ' cd ${CUR_DIR}\n'
657656
txt += 'done\n'
658657
txt += 'echo "Job ended at `date`"'
659658

@@ -771,6 +770,7 @@ def submit_slurm_array(
771770
array_max=None,
772771
dependency: Union[List[str],None] = None,
773772
group_size: int = 1,
773+
debug: bool = False,
774774
**kwargs,
775775
) -> Union[None,List[int]]:
776776
'''
@@ -802,8 +802,6 @@ def submit_slurm_array(
802802
nslurm_jobs = int(np.ceil(njobs / group_size))
803803
else:
804804
nslurm_jobs = njobs
805-
806-
# self._gen_array_script(group_size=group_size)
807805

808806
if dependency is not None:
809807
dependstr = None
@@ -826,6 +824,12 @@ def submit_slurm_array(
826824
'job_array.sh'
827825
]
828826

827+
if debug:
828+
# Only for testing purposes
829+
print('SLURM command:', command)
830+
os.environ['SLURM_ARRAY_TASK_ID'] = '0'
831+
command = ['sh', 'job_array.sh']
832+
829833
completed_process = subprocess.run(
830834
command, check=False, capture_output=True, text=True,
831835
)
@@ -842,7 +846,12 @@ def submit_slurm_array(
842846
logger.error(err_msg)
843847
completed_process.check_returncode()
844848

845-
job_ids = [int(completed_process.stdout.split(' ')[-1])]
849+
if debug:
850+
# logging.error('STDOUT:'+f'{completed_process.stdout}')
851+
logging.error('STDERR:'+f'{completed_process.stderr}')
852+
job_ids = None
853+
else:
854+
job_ids = [int(completed_process.stdout.split(' ')[-1])]
846855
return job_ids
847856

848857
def submit(self, **kwargs) -> None:
@@ -898,7 +907,7 @@ def get_last_output(self) -> Dict:
898907
''' Returns the output of the last job in the chain '''
899908
return self.unitjobs[-1].get_output()
900909

901-
def submit(self, dependency: Union[List,None] = None) -> List:
910+
def submit(self, dependency: Union[List,None] = None, debug: bool = False) -> List:
902911
'''
903912
Submit a job using slurm, interactively or in the terminal
904913
'''
@@ -997,7 +1006,7 @@ def submit(self, dependency: Union[List,None] = None) -> List:
9971006
if i == 0:
9981007
write_image = True
9991008

1000-
if only_write:
1009+
if only_write or debug:
10011010
curjob.gen_input_files(write_image=write_image)
10021011
else:
10031012
dependency = curjob.submit(
@@ -1031,29 +1040,26 @@ def submit(self, dependency: Union[List,None] = None) -> List:
10311040
return job_ids
10321041

10331042

1034-
def load_job_from_directory(workdir: str):
1043+
def load_job_from_directory(workdir: os.PathLike) -> Job:
10351044
''' Loads a job from a given directory '''
10361045
workdir = Path(workdir)
1037-
assert workdir.exists(), f'Work director {workdir} does not exist'
1046+
assert workdir.exists(), f'Work directory "{workdir}" does not exist'
10381047
logger = get_logger()
1039-
sim_inputs = glob(str(workdir / 'sim_input.yaml'))
1040-
if len(sim_inputs) != 1:
1041-
logger.error('Multiple or no sim_input.yaml files in %s', {workdir})
10421048
try:
1043-
sim_input = read_yaml(glob(str(workdir / 'sim_input.yaml'))[0])
1049+
sim_input = read_yaml(workdir / 'sim_input.yaml')
10441050
except IndexError as exc:
10451051
logger.error('sim_input.yaml not found in %s', {str(workdir)})
10461052
raise exc
10471053

1048-
env_inputs = glob(str(workdir / 'env_input.yaml'))
1049-
if len(env_inputs) == 1:
1050-
env_input = read_yaml(env_inputs[0])
1054+
env_input_file = workdir / 'env_input.yaml'
1055+
if env_input_file.exists():
1056+
env_input = read_yaml(env_input_file)
10511057
else:
10521058
env_input = None
10531059

1054-
calc_inputs = glob(str(workdir / 'calc_input.yaml'))
1055-
if len(calc_inputs) == 1:
1056-
calc_input = read_yaml(calc_inputs[0])
1060+
calc_input_file = workdir / 'calc_input.yaml'
1061+
if calc_input_file.exists():
1062+
calc_input = read_yaml(calc_input_file)
10571063
else:
10581064
calc_input = None
10591065

‎tests/asimmodules/workflows/test_distributed.py

+56-29
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,28 @@
22
Tests for running asimmodules using asim_run.py
33
"""
44
from glob import glob
5+
import os
56
import pytest
6-
from asimtools.job import create_unitjob
7-
from asimtools.job import load_job_from_directory
7+
from asimtools.job import (
8+
create_unitjob,
9+
DistributedJob,
10+
load_job_from_directory,
11+
)
12+
13+
def create_distjob(sim_input, env_input, workdir, calc_input=None):
14+
"""Helper for making a generic DistributedJob object, mostly for testing"""
15+
env_id = list(env_input.keys())[0]
16+
sim_input['env_id'] = env_id
17+
if calc_input is not None:
18+
calc_id = list(calc_input.keys())[0]
19+
sim_input['calc_id'] = calc_id
20+
sim_input['workdir'] = workdir
21+
distjob = DistributedJob(
22+
sim_input['args']['subsim_inputs'],
23+
env_input,
24+
calc_input,
25+
)
26+
return distjob
827

928
@pytest.mark.parametrize("calc_input",["lj_argon_calc_input"])
1029
@pytest.mark.parametrize("env_input",["inline_env_input"])
@@ -31,30 +50,38 @@ def test_distributed(env_input, calc_input, sim_input, tmp_path, request):
3150

3251
assert unitjob.get_status(descend=True) == (True, 'complete')
3352

34-
# @pytest.mark.parametrize("calc_input",["lj_argon_calc_input"])
35-
# @pytest.mark.parametrize("env_input",["batch_env_input"])
36-
# @pytest.mark.parametrize("sim_input",[
37-
# "lj_distributed_batch_sim_input",
38-
# "lj_distributed_group_batch_sim_input",
39-
# ])
40-
# def test_batch_distributed(env_input, calc_input, sim_input, tmp_path, request):
41-
# env_input = request.getfixturevalue(env_input)
42-
# calc_input = request.getfixturevalue(calc_input)
43-
# sim_input = request.getfixturevalue(sim_input)
44-
# wdir = tmp_path / 'wdir'
45-
# unitjob = create_unitjob(sim_input, env_input, wdir, calc_input=calc_input)
46-
# unitjob.submit()
47-
48-
# assert load_job_from_directory(wdir).get_status()[1] == 'complete'
49-
# dirs = glob(str(wdir / 'id*'))
50-
# assert len(dirs) == len(sim_input['args']['subsim_inputs'])
51-
52-
# for d in dirs:
53-
# assert str(d).rsplit('/', maxsplit=1)[-1].startswith('id-')
54-
55-
# uj = load_job_from_directory(d)
56-
# assert uj.get_status()[1] == 'complete'
57-
58-
# assert uj.get_sim_input()['workdir'] == './'
59-
60-
# assert unitjob.get_status(descend=True) == (True, 'complete')
53+
@pytest.mark.parametrize("calc_input",["lj_argon_calc_input"])
54+
@pytest.mark.parametrize("env_input",["batch_env_input"])
55+
@pytest.mark.parametrize("sim_input",[
56+
"lj_distributed_batch_sim_input",
57+
"lj_distributed_group_batch_sim_input",
58+
])
59+
def test_batch_distributed(env_input, calc_input, sim_input, tmp_path, request):
60+
env_input = request.getfixturevalue(env_input)
61+
calc_input = request.getfixturevalue(calc_input)
62+
sim_input = request.getfixturevalue(sim_input)
63+
group_size = sim_input['args'].get('group_size', 1)
64+
wdir = tmp_path
65+
os.chdir(wdir)
66+
distjob = create_distjob(sim_input, env_input, wdir, calc_input=calc_input)
67+
group_size=sim_input['args'].get('group_size', 1)
68+
array_max=sim_input['args'].get('array_max', None)
69+
skip_failed=sim_input['args'].get('skip_failed', False)
70+
distjob.submit(
71+
debug=True,
72+
group_size=group_size,
73+
array_max=array_max,
74+
skip_failed=skip_failed,
75+
)
76+
dirs = glob(str(wdir / 'id*'))
77+
assert len(dirs) == len(sim_input['args']['subsim_inputs'])
78+
79+
statuses = ['complete'] * group_size + ['clean'] * (len(dirs) - 1)
80+
for d_ind, d in enumerate(dirs):
81+
assert str(d).rsplit('/', maxsplit=1)[-1].startswith('id-')
82+
83+
uj = load_job_from_directory(d)
84+
print(uj.workdir, uj.get_status())
85+
assert uj.get_status()[1] == statuses[d_ind]
86+
87+
# assert distjob.get_status(descend=False) == (True, 'complete')

‎tests/conftest.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -290,7 +290,7 @@ def lj_distributed_batch_sim_input():
290290
'''
291291
subsim_input = {
292292
'asimmodule': 'singlepoint',
293-
'env_id': 'inline',
293+
'env_id': 'batch',
294294
'args': {
295295
'calc_id': 'lj',
296296
'image': {
@@ -321,7 +321,7 @@ def lj_distributed_group_batch_sim_input():
321321
'''
322322
subsim_input = {
323323
'asimmodule': 'singlepoint',
324-
'env_id': 'inline', # This should be overwrriten by the group env
324+
'env_id': 'batch', # This should be overwrriten by the group env
325325
'args': {
326326
'calc_id': 'lj',
327327
'image': {
@@ -346,3 +346,4 @@ def lj_distributed_group_batch_sim_input():
346346
}
347347

348348
return sim_input
349+

0 commit comments

Comments
 (0)