Coverage for  / opt / hostedtoolcache / Python / 3.10.19 / x64 / lib / python3.10 / site-packages / starlord / cli.py: 66%

1import argparse 

2import pathlib 

3import sys 

4 

5from . import __version__ 

6from .grid_gen import GridGenerator 

7from .star_fitter import StarFitter 

8 

9if sys.version_info >= (3, 11): 

10 import tomllib 

11else: 

12 import tomli as tomllib 

13 

14 

15def main(): 

16 parser = argparse.ArgumentParser( 

17 "starlord", description="Fit stellar observations with starlord from the command line.") 

18 parser.add_argument( 

19 "input", type=pathlib.Path, nargs="?", default=None, help="A toml file to load run settings from (optional)") 

20 parser.add_argument("--dry-run", action="store_true") 

21 parser.add_argument("-v", "--verbose", action="store_true") 

22 parser.add_argument("-c", "--code", action="store_true", help="Print code upon generation.") 

23 parser.add_argument("--version", action="version", version=f"starlord {__version__}") 

24 parser.add_argument("-l", "--list-grids", action="store_true") 

25 args = parser.parse_args(args=None if sys.argv[1:] else ['--help']) 

26 

27 if args.list_grids: 

28 if args.input is not None: 

29 grid_name = str(args.input) 

30 assert grid_name in GridGenerator.grids(), f"Grid {grid_name} not found." 

31 g = GridGenerator.get_grid(grid_name) 

32 print("Grid:", g.name) 

33 print("Inputs:", ", ".join(g.inputs)) 

34 print("Outputs:") 

35 print(*[" " + i for i in g.outputs], sep="\n") 

36 if len(g.derived) > 0: 

37 print("Derived:") 

38 print(*g.derived, sep="\n") 

39 return 

40 print("Available grids:") 

41 for g in GridGenerator.grids().values(): 

42 # Print short grid info, no need for "Grid_" prefix. 

43 print(" ", str(g)[5:]) 

44 return 

45 

46 # === Load Settings === 

47 # Default initial settings (keep minimal) 

48 settings = {'output': {'terminal': True, 'file': ""}, "sampling": {}} 

49 

50 if args.input is not None: 

51 if args.verbose: 

52 print(args.input) 

53 with open(args.input, 'rb') as f: 

54 settings.update(tomllib.load(f)) 

55 # TODO: Handle syntax errors in the toml file 

56 

57 # Report ignored sections 

58 for section in settings.keys(): 

59 if section not in ['model', 'sampling', 'output']: 

60 print(f"Warning, section {section} in input file {args.input} is not used.") 

61 

62 # TODO: Update settings with command line arguments 

63 

64 if args.verbose: 

65 print("Args:", args) 

66 print("Settings: ", settings) 

67 

68 # === Setup the fitter === 

69 assert "model" in settings.keys(), "No model information was specified." 

70 fitter = StarFitter(args.verbose) 

71 fitter.set_from_dict(settings['model']) 

72 if args.code: 

73 # TODO: Set prior type based on sampler type 

74 print(fitter.generate()) 

75 if args.dry_run: 

76 # TODO: Check constants 

77 fitter.summary(args.verbose) 

78 return 

79 

80 # === Run Sampler == 

81 consts = settings['sampling'].get('const', {}) 

82 if args.verbose: 

83 print("Constants:", consts) 

84 results = fitter.run_sampler({}, constants=consts) 

85 

86 # === Write Outputs === 

87 out: dict = {"terminal": False, "file": ""} 

88 out.update(settings['output']) 

89 if out['terminal']: 

90 print(results.summary()) 

91 if out['file'] != "": 

92 print("TODO: write results to ", out['file']) 

93 # fitter.write_results()