updates

src/methane_super_emitters/optimize.py

@@ -6,8 +6,9 @@
 
 @click.command()
 @click.option("-i", "--input-dir", help="Data directory")
-@click.option("-m", "--max-epochs", help="Maximum number of epochs", default=1)
-def train_model(input_dir, max_epochs):
+@click.option("-m", "--max-epochs", help="Maximum number of epochs", default=100)
+@click.option("-n", "--n-trials", help="Number of trials or points to sample", default=200)
+def optimize_model(input_dir, max_epochs, n_trials):
     def objective(trial):
         fields = ["methane", "u10", "v10", "qa"]
         dropout_rate = trial.suggest_float("dropout", 0.1, 0.9)
@@ -18,10 +19,10 @@ def objective(trial):
         trainer.fit(model=model, datamodule=datamodule)
         return trainer.callback_metrics['val_acc'].item()
     study = optuna.create_study(direction='maximize')
-    study.optimize(objective, n_trials=200)
+    study.optimize(objective, n_trials=n_trials)
     df = study.trials_dataframe()
     df.to_csv('opt_results.csv')
     print("Best parameters:", study.best_params)
 
 if __name__ == "__main__":
-    train_model()
+    optimize_model()

test/test_main.py

@@ -2,10 +2,14 @@
 import lightning as L
 from methane_super_emitters.model import SuperEmitterDetector
 from methane_super_emitters.datamodule import TROPOMISuperEmitterDataModule
+from methane_super_emitters.optimize import optimize_model
 
 def test_model():
     model = SuperEmitterDetector(fields=['methane', 'qa', 'u10', 'v10'])
     datamodule = TROPOMISuperEmitterDataModule('./data/dataset', fields=['methane', 'qa', 'u10', 'v10'])
     trainer = L.Trainer(max_epochs=10)
     trainer.fit(model=model, datamodule=datamodule)
     trainer.test(model=model, datamodule=datamodule)
+
+def test_optimize():
+    optimize_model('data/dataset', 1, 10)