google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
#!/usr/bin/python
17
r"""Generates every neutral stoichiometry with a given number of heavy atoms.
18

19
Example usage:
20
mkdir stoichs
21
enumerate_stoichiometries.py --output_prefix=stoichs/ \
22
    --num_heavy=3 --heavy_elements=C,N,O,S
23
"""
24

25
from absl import app
26
from absl import flags
27

28
from graph_sampler import stoichiometry
29

30
FLAGS = flags.FLAGS
31

32
flags.DEFINE_integer('num_heavy', None, 'Number of non-hydrogen atoms.')
33
flags.DEFINE_list('heavy_elements', ['C', 'N', 'N+', 'O', 'O-', 'F'],
34
                  'Which heavy elements to use.')
35
flags.DEFINE_string('output_prefix', '', 'Prefix for output files.')
36
flags.DEFINE_list(
37
    'valences', [],
38
    'Valences of atom types (only required for atom types whose valence cannot '
39
    'be inferred by rdkit, (e.g. "X=7,R=3" if you\'re using "synthetic atoms" '
40
    'with valences 7 and 3).')
41
flags.DEFINE_list(
42
    'charges', [],
43
    'Charges of atom types (only required for atom types whose charge cannot '
44
    'be inferred by rdkit, (e.g. "X=0,R=-1" if you\'re using "synthetic atoms" '
45
    'with valences 0 and -1).')
46

47

48
def main(argv):
49
  if len(argv) > 1:
50
    raise RuntimeError(f'Unexpected arguments: {argv[1:]}')
51
  FLAGS.valences = stoichiometry.parse_dict_flag(FLAGS.valences)
52
  FLAGS.charges = stoichiometry.parse_dict_flag(FLAGS.charges)
53

54
  count = 0
55
  for stoich in stoichiometry.enumerate_stoichiometries(FLAGS.num_heavy,
56
                                                        FLAGS.heavy_elements,
57
                                                        FLAGS.valences,
58
                                                        FLAGS.charges):
59
    element_str = ''.join(stoich.to_element_list())
60
    fn = '%s%d_%s.stoich' % (FLAGS.output_prefix, FLAGS.num_heavy, element_str)
61
    print(element_str)
62
    with open(fn, 'w') as f:
63
      stoich.write(f)
64
    count += 1
65

66
  print(f'{count} files written!')
67

68

69
if __name__ == '__main__':
70
  flags.mark_flag_as_required('num_heavy')
71
  app.run(main)
72