remove module when importing trait extensions and using runner with c…

…ustom args

src/contact.rs

@@ -1,9 +1,9 @@
 use ixa::{
-    context::Context,
+    Context,
     define_rng,
-    error::IxaError,
-    people::{ContextPeopleExt, PersonId},
-    random::ContextRandomExt,
+    IxaError,
+    ContextPeopleExt, PersonId,
+    ContextRandomExt,
 };
 
 use crate::population_loader::Alive;
@@ -54,7 +54,7 @@ impl ContextContactExt for Context {
 mod test {
     use super::ContextContactExt;
     use crate::population_loader::Alive;
-    use ixa::{context::Context, people::ContextPeopleExt, random::ContextRandomExt, IxaError};
+    use ixa::{Context, ContextPeopleExt, ContextRandomExt, IxaError};
 
     #[test]
     fn test_cant_get_contact_in_pop_of_one() {

src/main.rs

@@ -3,50 +3,55 @@ mod parameters;
 mod population_loader;
 mod transmission_manager;
 
-use clap::Parser;
+use ixa::runner::run_with_custom_args;
+use clap::Args;
 use ixa::{
-    context::Context, error::IxaError, global_properties::ContextGlobalPropertiesExt,
-    people::ContextPeopleExt, random::ContextRandomExt, report::ContextReportExt,
+    Context, IxaError, ContextGlobalPropertiesExt,
+    ContextPeopleExt, ContextRandomExt, ContextReportExt,
 };
 use std::path::PathBuf;
 use transmission_manager::InfectiousStatus;
 
 use crate::parameters::Parameters;
 use crate::population_loader::{Age, CensusTract};
 
-#[derive(Parser)]
-#[command(version, about, long_about = None)]
-struct Args {
-    /// path to the input file
-    #[arg(short, long)]
-    input_file: PathBuf,
-
-    /// path to the output directory
-    #[arg(short, long)]
-    output_directory: PathBuf,
-
+#[derive(Args, Debug)]
+struct CustomArgs {
     /// whether force overwrite of output files if they already exist
-    #[arg(short, long, default_value = "false", default_missing_value = "true")]
+    #[arg(short = 'f', long)]
     force_overwrite: bool,
 }
 
-fn initialize(args: &Args) -> Result<Context, IxaError> {
-    let mut context = Context::new();
-    // Read the global properties.
-    context.load_global_properties(&args.input_file)?;
+fn initialize() -> Result<Context, IxaError> {
+    let mut context = run_with_custom_args(|context, args, custom_args: Option<CustomArgs>| {
+        println!("Setting Overwrite parameter");
+        // Read the global properties.
+        if let Some(custom_args)  = custom_args {            
+            // Set the output directory and whether to overwrite the existing file.
+            context
+                .report_options()
+                .directory(PathBuf::from(&args.output_dir))
+                .overwrite(custom_args.force_overwrite);
+        } else {
+            // If overwrite not specified, set to false (default)
+            context
+                .report_options()
+                .directory(PathBuf::from(&args.output_dir));
+        }
+
+        Ok(())
+    })
+    .unwrap();
+
+
     let parameters = context
         .get_global_property_value(Parameters)
         .unwrap()
         .clone();
+
     // Set the random seed.
     context.init_random(parameters.seed);
 
-    // Set the output directory and whether to overwrite the existing file.
-    context
-        .report_options()
-        .directory(PathBuf::from(&args.output_directory))
-        .overwrite(args.force_overwrite);
-
     // Report the number of people by age, census tract, and infectious status
     // every report_period.
     context.add_periodic_report(
@@ -62,7 +67,7 @@ fn initialize(args: &Args) -> Result<Context, IxaError> {
     context.index_property(CensusTract);
 
     // Initialize the person-to-person transmission workflow.
-    transmission_manager::init(&mut context);
+    transmission_manager::init(&mut context)?;
 
     // Add a plan to shut down the simulation after `max_time`, regardless of
     // what else is happening in the model.
@@ -78,7 +83,6 @@ fn initialize(args: &Args) -> Result<Context, IxaError> {
 }
 
 fn main() {
-    let args = Args::parse();
-    let mut context = initialize(&args).expect("Error initializing.");
+    let mut context = initialize().expect("Error initializing.");
     context.execute();
 }

src/parameters.rs

@@ -1,7 +1,7 @@
 use std::fmt::Debug;
 use std::path::PathBuf;
 
-use ixa::{define_global_property, error::IxaError};
+use ixa::{define_global_property, IxaError};
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -33,7 +33,7 @@ define_global_property!(Parameters, ParametersValues, validate_inputs);
 
 #[cfg(test)]
 mod test {
-    use ixa::error::IxaError;
+    use ixa::IxaError;
 
     use super::validate_inputs;
     use std::path::PathBuf;

src/population_loader.rs

@@ -1,7 +1,8 @@
 use crate::parameters::Parameters;
 use ixa::{
-    context::Context, define_person_property, define_person_property_with_default, error::IxaError,
-    global_properties::ContextGlobalPropertiesExt, people::ContextPeopleExt,
+    Context, define_person_property, define_person_property_with_default,
+    IxaError,
+    ContextGlobalPropertiesExt, ContextPeopleExt,
 };
 
 use serde::Deserialize;
@@ -56,7 +57,7 @@ pub fn init(context: &mut Context) -> Result<(), IxaError> {
 #[cfg(test)]
 mod test {
     use super::*;
-    use ixa::people::ContextPeopleExt;
+    use ixa::ContextPeopleExt;
     use std::io::Write;
     use std::path::PathBuf;
     use tempfile::NamedTempFile;

src/transmission_manager.rs

@@ -1,15 +1,15 @@
 use ixa::{
-    context::Context,
+    Context,
     define_person_property, define_person_property_with_default, define_rng,
-    error::IxaError,
-    global_properties::ContextGlobalPropertiesExt,
-    people::{ContextPeopleExt, PersonId, PersonPropertyChangeEvent},
-    random::ContextRandomExt,
+    IxaError,
+    ContextGlobalPropertiesExt,
+    ContextPeopleExt, PersonId, PersonPropertyChangeEvent,
+    ContextRandomExt,
 };
 use statrs::distribution::{ContinuousCDF, Exp, Poisson};
 
 use crate::parameters::Parameters;
-use crate::{contact::ContextContactExt, population_loader::Alive};
+use crate::contact::ContextContactExt;
 
 // Define the possible infectious statuses for a person.
 // These states refer to the person's infectiousness at a given time
@@ -31,31 +31,32 @@ define_person_property_with_default!(
     InfectiousStatusType::Susceptible
 );
 
-/// Seeds initial infections at t = 0, and subscribes to people becoming infectious
-/// to schedule their infection attempts.
-pub fn init(context: &mut Context) {
+/// Seeds initial infections at t = 0, and subscribes to
+/// people becoming infectious to schedule their infection attempts.
+pub fn init(context: &mut Context) -> Result<(), IxaError> {
     // Watch for changes in the InfectiousStatusType property.
     context.subscribe_to_event(
         move |context, event: PersonPropertyChangeEvent<InfectiousStatus>| {
             handle_infectious_status_change(context, event);
         },
     );
     context.add_plan(0.0, |context| {
-        seed_infections(context);
+        seed_infections(context).expect("Unable to seed infections");
     });
+    Ok(())
 }
 
 /// This function seeds the initial infections in the population.
-fn seed_infections(context: &mut Context) {
+fn seed_infections(context: &mut Context) -> Result<(), IxaError> {
     // For now, we just pick a random person and make them infectious.
     // In the future, we may pick people based on specific person properties.
-    let alive_people = context.query_people((Alive, true));
-    let random_person_id = context.sample_range(TransmissionRng, 0..alive_people.len());
+    let random_person_id = context.sample_person(TransmissionRng)?;
     context.set_person_property(
-        alive_people[random_person_id],
+        random_person_id,
         InfectiousStatus,
         InfectiousStatusType::Infectious,
     );
+    Ok(())
 }
 
 // Called when a person's infectious status changes. Only considers people becoming infectious,
@@ -227,8 +228,8 @@ mod test {
 
     use super::{init, InfectiousStatus, InfectiousStatusType};
     use ixa::{
-        context::Context, global_properties::ContextGlobalPropertiesExt, people::ContextPeopleExt,
-        random::ContextRandomExt, PersonId, PersonPropertyChangeEvent,
+        Context, ContextGlobalPropertiesExt, ContextPeopleExt,
+        ContextRandomExt, PersonId, PersonPropertyChangeEvent,
     };
     use statrs::distribution::{ContinuousCDF, Exp};
 
@@ -256,7 +257,7 @@ mod test {
         // and we do not trigger the `get_contact` function -- which errors out in the case
         // of a population size of 1.
         let mut context = setup(0.000_000_000_000_000_01);
-        init(&mut context);
+        init(&mut context).expect("could not initialize");
         let person_id = context.add_person(()).unwrap();
 
         context.execute();
@@ -273,7 +274,7 @@ mod test {
         // zero secondary infections is extremely low.
         // This lets us check that the other person in the population is infected.
         let mut context = setup(50.0);
-        init(&mut context);
+        init(&mut context).expect("could not initialize");
         let person_id = context.add_person(()).unwrap();
         let contact = context.add_person(()).unwrap();
 
@@ -339,7 +340,7 @@ mod test {
             // By having only one person in the population when we choose the transmitter, we guarantee
             // that that one person is the transmitter. This lets us keep the transmitter id, which we
             // need below.
-            init(&mut context);
+            init(&mut context).expect("could not initialize");
             let infection_times_clone = Rc::clone(&infection_times);
             context.subscribe_to_event({
                 move |context, event: PersonPropertyChangeEvent<InfectiousStatus>| {