File size: 1,510 Bytes
35cdf53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
"""Script for building intermediate data."""
import os
from importlib import resources
import pathlib
import site
import flax_model.alphafold3.constants.converters
from flax_model.alphafold3.constants.converters import ccd_pickle_gen
from flax_model.alphafold3.constants.converters import chemical_component_sets_gen
def build_data():
"""Builds intermediate data."""
explicit_cif_path = os.environ.get('ALPHAFOLD3_CIFPP_COMPONENTS')
if explicit_cif_path:
cif_path = pathlib.Path(explicit_cif_path)
if not cif_path.exists():
raise ValueError(f'Configured components.cif does not exist: {cif_path}')
else:
for site_path in site.getsitepackages():
path = pathlib.Path(site_path) / 'share/libcifpp/components.cif'
if path.exists():
cif_path = path
break
else:
raise ValueError('Could not find components.cif')
output_root = os.environ.get('ALPHAFOLD3_DATA_OUTPUT_DIR')
if output_root:
out_root = pathlib.Path(output_root)
out_root.mkdir(parents=True, exist_ok=True)
else:
out_root = pathlib.Path(resources.files(flax_model.alphafold3.constants.converters))
ccd_pickle_path = out_root / 'ccd.pickle'
chemical_component_sets_pickle_path = out_root / 'chemical_component_sets.pickle'
ccd_pickle_gen.main(['', str(cif_path), str(ccd_pickle_path)])
chemical_component_sets_gen.main(
['', str(chemical_component_sets_pickle_path)]
)
def main() -> None:
build_data()
if __name__ == "__main__":
main()
|