def get_model_spec_time_varying_with_covariates(covariate_delay):
    """Return model specification with hyperparameters and covariates."""
    result = model_spec_lib.ModelSpec(
        hparams={
            "initial_learning_rate": 0.003061242915556316,
            "momentum": 0.2,
            "decay_steps": 1000,
            "fine_tuning_steps": 100,
            "fine_tuning_decay": 1.0,
            "decay_rate": 1.0,
            "location_dependent_init": False,
            "infected_threshold": 3,
            "restart_threshold": 1000,
            "time_scale_weight": 0.00006243159539906051,
            "train_loss_coefs":
            [0, 0.001, 0.2, 0.1, 0.01, 0.01, 0.2, 0.1, 0.01],
            "valid_loss_coefs":
            [0, 0.001, 0.2, 0.1, 0.01, 0.01, 0.2, 0.1, 0.01],
            "sync_coef": 1.0,
            "reduced_sync_undoc": 1.0,
            "smooth_coef": 0.5,
            "first_dose_vaccine_ratio_per_day_init": 0.0,
            "second_dose_vaccine_ratio_per_day_init": 0.0,
            "average_contact_id_rate_init": -1.9131825459930378,
            "average_contact_iud_rate_init": -1.071866945725303,
            "reinfectable_rate_init": -5.548940468292865,
            "alpha_rate_init": -2.272765554778715,
            "diagnosis_rate_init": -2.095597433974376,
            "recovery_id_rate_init": -1.495660223962899,
            "recovery_iud_rate_init": -1.475605314236803,
            "recovery_h_rate_init": -1.9032896753850963,
            "hospitalization_rate_init": -1.4331763640928012,
            "death_id_rate_init": -1.8060447968974489,
            "death_h_rate_init": -1.4886876719378206,
            "bias_penalty_coef": 0.2835406167308398,
            "r_eff_penalty_coef": 2.0,
            "acceleration_death_coef": 0.1,
            "acceleration_confirm_coef": 0.1,
            "acceleration_hospital_coef": 0.1,
            "quantile_encoding_window": 7,
            "quantile_smooth_coef": 0.9,
            "quantile_training_iteration_ratio": 0.5,
            "width_coef_train": 10.0,
            "width_coef_valid": 5.0,
            "quantile_cum_viol_coef": 500.0,
            "increment_loss_weight": 0.0,
            "lasso_penalty_coef": 1.0,
            "covariate_training_mixing_coef": 1.0,
            "train_window_range": 2.0,
            "partial_mean_interval": 4,
            "direction_loss_coef": 2000.0,
            "train_crps_weight": 0.25,
            "valid_crps_weight": 0.25
        },
        encoder_specs=[
            model_spec_lib.EncoderSpec(
                encoder_name="first_dose_vaccine_ratio_per_day",
                encoder_type="vaccine",
                vaccine_type="first_dose",
                covariate_feature_specs=[
                    first_dose_vaccine_ratio_per_day,
                    first_dose_vaccine_effectiveness,
                    second_dose_vaccine_effectiveness
                ],
            ),
            model_spec_lib.EncoderSpec(
                encoder_name="second_dose_vaccine_ratio_per_day",
                encoder_type="vaccine",
                vaccine_type="second_dose",
                covariate_feature_specs=[
                    second_dose_vaccine_ratio_per_day,
                    first_dose_vaccine_effectiveness,
                    second_dose_vaccine_effectiveness
                ],
            ),
            model_spec_lib.EncoderSpec(
                encoder_name="average_contact_id_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km,
                    income_per_capita,
                    aqi_mean_2018,
                    household,
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed, preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum, mobility_index,
                    mobility_samples, confirmed_per_csrp_tests
                ] + extra_covariates + npi_features,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="average_contact_iud_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km,
                    income_per_capita,
                    aqi_mean_2018,
                    household,
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed, preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum, mobility_index,
                    mobility_samples, confirmed_per_csrp_tests
                ] + extra_covariates + npi_features,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="reinfectable_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km, household,
                    household_food_stamp, income_per_capita,
                    patient_experience_above_the_national_average,
                    patient_experience_below_the_national_average,
                    patient_experience_same_as_the_national_average,
                    rating_average, aqi_mean_2018, kaiser_60plus_population
                ] + extra_static_features,
                covariate_feature_specs=None,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                }),
            model_spec_lib.EncoderSpec(encoder_name="alpha_rate",
                                       encoder_type="gam",
                                       static_feature_specs=[],
                                       covariate_feature_specs=None,
                                       encoder_kwargs={
                                           "initial_bias": 0,
                                           "location_dependent_bias": True,
                                       }),
            model_spec_lib.EncoderSpec(
                encoder_name="diagnosis_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km, household,
                    household_food_stamp, income_per_capita,
                    patient_experience_above_the_national_average,
                    patient_experience_below_the_national_average,
                    patient_experience_same_as_the_national_average,
                    rating_average, aqi_mean_2018, kaiser_60plus_population
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed,
                    preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    confirmed_per_csrp_tests,
                    total_tests_per_capita,
                    dow_feature,
                ] + extra_covariates,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="recovery_id_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km, household,
                    household_food_stamp, income_per_capita,
                    patient_experience_above_the_national_average,
                    patient_experience_below_the_national_average,
                    patient_experience_same_as_the_national_average,
                    rating_average, aqi_mean_2018, kaiser_60plus_population
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed,
                    preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    dow_feature,
                ] + extra_covariates,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="recovery_iud_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km, household,
                    household_food_stamp, income_per_capita,
                    patient_experience_above_the_national_average,
                    patient_experience_below_the_national_average,
                    patient_experience_same_as_the_national_average,
                    rating_average, aqi_mean_2018, kaiser_60plus_population
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed, preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum
                ] + extra_covariates,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="recovery_h_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km, household,
                    household_food_stamp, income_per_capita,
                    patient_experience_above_the_national_average,
                    patient_experience_below_the_national_average,
                    patient_experience_same_as_the_national_average,
                    rating_average, aqi_mean_2018, kaiser_60plus_population,
                    icu_beds
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed,
                    preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    dow_feature,
                ] + extra_covariates,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="hospitalization_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km, household,
                    household_food_stamp, income_per_capita,
                    patient_experience_above_the_national_average,
                    patient_experience_below_the_national_average,
                    patient_experience_same_as_the_national_average,
                    rating_average, aqi_mean_2018, kaiser_60plus_population
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed,
                    preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    dow_feature,
                    confirmed_per_csrp_tests,
                    total_tests_per_capita,
                ] + extra_covariates,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="death_id_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km, household,
                    household_food_stamp, income_per_capita,
                    patient_experience_above_the_national_average,
                    patient_experience_below_the_national_average,
                    patient_experience_same_as_the_national_average,
                    rating_average, aqi_mean_2018, kaiser_60plus_population,
                    icu_beds
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed,
                    preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    dow_feature,
                ] + extra_covariates,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="death_h_rate",
                encoder_type="gam",
                static_feature_specs=[
                    population_density_per_sq_km, household,
                    household_food_stamp, income_per_capita,
                    patient_experience_above_the_national_average,
                    patient_experience_below_the_national_average,
                    patient_experience_same_as_the_national_average,
                    rating_average, aqi_mean_2018, kaiser_60plus_population,
                    icu_beds
                ] + extra_static_features,
                covariate_feature_specs=[
                    preprocessed_confirmed,
                    preprocessed_death,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    dow_feature,
                ] + extra_covariates,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
        ])
    return result
     "bias_penalty_coef": 0.27315803186400106,
     "boundary_penalty_coef": 2.600265396382176,
     "lasso_penalty_coef": 1.0,
     "r_eff_penalty_coef": 1.0,
     "quantile_encoding_window": 7,
     "quantile_smooth_coef": 0.8,
     "increment_loss_weight": 0.0,
     "covariate_training_mixing_coef": 1.0,
     "direction_loss_coef": 2000.0
 },
 encoder_specs=[
     model_spec_lib.EncoderSpec(
         encoder_name="first_dose_vaccine_ratio_per_day",
         encoder_type="vaccine",
         vaccine_type="first_dose",
         covariate_feature_specs=[
             first_dose_vaccine_ratio_per_day,
             first_dose_vaccine_effectiveness,
             second_dose_vaccine_effectiveness
         ],
     ),
     model_spec_lib.EncoderSpec(
         encoder_name="second_dose_vaccine_ratio_per_day",
         encoder_type="vaccine",
         vaccine_type="second_dose",
         covariate_feature_specs=[
             second_dose_vaccine_ratio_per_day,
             first_dose_vaccine_effectiveness,
             second_dose_vaccine_effectiveness
         ],
     ),
     model_spec_lib.EncoderSpec(
def get_model_spec_time_varying_with_covariates(covariate_delay):
  """Return model specification with hyperparameters and covariates."""
  result = model_spec_lib.ModelSpec(
      hparams={
          "initial_learning_rate": 0.003778327764151733,
          "momentum": 0.1,
          "decay_steps": 1000,
          "fine_tuning_steps": 100,
          "fine_tuning_decay": 1.0,
          "decay_rate": 1.0,
          "location_dependent_init": False,
          "infected_threshold": 10,
          "restart_threshold": 1000,
          "time_scale_weight": 0.00006243159539906051,
          "train_loss_coefs": [0, 0.001, 1.0, 0.1, 0.05, 0.05, 0.01, 0.001],
          "valid_loss_coefs": [0, 0.001, 1.0, 0.1, 0.05, 0.05, 0.01, 0.001],
          "sync_coef": 1.0,
          "reduced_sync_undoc": 1.0,
          "smooth_coef": 0.5,
          "first_dose_vaccine_ratio_per_day_init": 0.0,
          "second_dose_vaccine_ratio_per_day_init": 0.0,
          "average_contact_id_rate_init": -1.9131825459930378,
          "average_contact_iud_rate_init": -1.071866945725303,
          "reinfectable_rate_init": -5.548940468292865,
          "alpha_rate_init": -2.272765554778715,
          "diagnosis_rate_init": -2.095597433974376,
          "recovery_id_rate_init": -1.495660223962899,
          "recovery_iud_rate_init": -1.475605314236803,
          "recovery_h_rate_init": -1.9032896753850963,
          "recovery_i_rate_init": -1.354652366086861,
          "recovery_v_rate_init": -1.2872596880548643,
          "hospitalization_rate_init": -1.4331763640928012,
          "icu_rate_init": -1.7697853537123167,
          "ventilator_rate_init": -1.747746088542831,
          "death_id_rate_init": -1.8060447968974489,
          "death_h_rate_init": -1.4886876719378206,
          "death_i_rate_init": -1.525990911868795,
          "death_v_rate_init": -1.0072190496934017,
          "bias_penalty_coef": 0.2835406167308398,
          "r_eff_penalty_coef": 2.0,
          "r_eff_penalty_cutoff": 5.5,
          "acceleration_death_coef": 0.1,
          "acceleration_confirm_coef": 0.1,
          "acceleration_hospital_coef": 0.1,
          "quantile_encoding_window": 7,
          "quantile_smooth_coef": 0.9,
          "quantile_training_iteration_ratio": 1.0,
          "width_coef_train": 1.0,
          "width_coef_valid": 1.0,
          "quantile_cum_viol_coef": 1000.0,
          "increment_loss_weight": 0.0,
          "lasso_penalty_coef": 1.0,
          "covariate_training_mixing_coef": 1.0,
          "direction_loss_coef": 15000.0,
          "train_window_range": 2.0,
          "partial_mean_interval": 4,
          "train_crps_weight": 0.5,
          "valid_crps_weight": 0.5
      },
      encoder_specs=[
          model_spec_lib.EncoderSpec(
              encoder_name="first_dose_vaccine_ratio_per_day",
              encoder_type="vaccine",
              vaccine_type="first_dose",
              covariate_feature_specs=[
                  first_dose_vaccine_ratio_per_day,
                  first_dose_vaccine_effectiveness,
                  second_dose_vaccine_effectiveness
              ],
          ),
          model_spec_lib.EncoderSpec(
              encoder_name="second_dose_vaccine_ratio_per_day",
              encoder_type="vaccine",
              vaccine_type="second_dose",
              covariate_feature_specs=[
                  second_dose_vaccine_ratio_per_day,
                  first_dose_vaccine_effectiveness,
                  second_dose_vaccine_effectiveness
              ],
          ),
          model_spec_lib.EncoderSpec(
              encoder_name="average_contact_id_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, income_per_capita,
                  aqi_mean_2018, household, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed, preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum, mobility_index,
                  mobility_samples, average_temperature, rainfall, snowfall,
                  antigen_positive_ratio, antibody_negative_ratio,
                  confirmed_per_total_tests
              ] + npi_features,
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="average_contact_iud_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, income_per_capita,
                  aqi_mean_2018, household, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed, preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum, mobility_index,
                  mobility_samples, average_temperature, rainfall, snowfall,
                  antigen_positive_ratio, antibody_negative_ratio,
                  confirmed_per_total_tests
              ] + npi_features,
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="reinfectable_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, income_per_capita,
                  household_food_stamp, household,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  aqi_mean_2018, population_60plus_ratio
              ],
              covariate_feature_specs=None,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="alpha_rate",
              encoder_type="gam",
              static_feature_specs=[],
              covariate_feature_specs=None,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="diagnosis_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, income_per_capita, household,
                  rating_average, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  total_tests_per_capita,
                  confirmed_per_total_tests,
                  dow_feature,
                  antigen_positive_ratio,
                  antibody_negative_ratio,
                  average_temperature,
                  rainfall,
                  snowfall,
              ] + symptoms_all,
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="recovery_id_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, income_per_capita,
                  household_food_stamp,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
                  antigen_positive_ratio,
                  antibody_negative_ratio,
                  average_temperature,
                  rainfall,
                  snowfall,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="recovery_iud_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, income_per_capita,
                  household_food_stamp,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  antigen_positive_ratio,
                  antibody_negative_ratio,
                  average_temperature,
                  rainfall,
                  snowfall,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="recovery_h_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, income_per_capita,
                  household_food_stamp,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="recovery_i_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, income_per_capita,
                  household_food_stamp,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="recovery_v_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, income_per_capita,
                  household_food_stamp,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="hospitalization_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km,
                  household,
                  household_food_stamp,
                  household,
                  income_per_capita,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services,
                  rating_average,
                  aqi_mean_2018,
                  icu_beds,
                  population_60plus_ratio,
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed, preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum, dow_feature,
                  total_tests_per_capita, confirmed_per_total_tests,
                  antigen_positive_ratio, antibody_negative_ratio,
                  average_temperature, rainfall, snowfall
              ] + symptoms_all,
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="icu_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, household_food_stamp,
                  household, income_per_capita,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  aqi_mean_2018, icu_beds, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="ventilator_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, household_food_stamp,
                  household, income_per_capita,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  aqi_mean_2018, icu_beds, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="death_id_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, household_food_stamp,
                  income_per_capita,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  aqi_mean_2018, icu_beds, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="death_h_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, household_food_stamp,
                  income_per_capita,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  aqi_mean_2018, icu_beds, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="death_i_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, household_food_stamp,
                  income_per_capita,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  aqi_mean_2018, icu_beds, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
          model_spec_lib.EncoderSpec(
              encoder_name="death_v_rate",
              encoder_type="gam",
              static_feature_specs=[
                  population_density_per_sq_km, household, household_food_stamp,
                  income_per_capita,
                  patient_experience_above_the_national_average,
                  patient_experience_below_the_national_average,
                  patient_experience_same_as_the_national_average,
                  hospital_type_critical_access_hospitals,
                  hospital_type_acute_care_department_of_defense,
                  service_type_emergency_services_supported,
                  service_type_non_emergency_services, rating_average,
                  aqi_mean_2018, icu_beds, population_60plus_ratio
              ],
              covariate_feature_specs=[
                  preprocessed_confirmed,
                  preprocessed_death,
                  preprocessed_confirmed_mean_to_sum,
                  preprocessed_death_mean_to_sum,
                  dow_feature,
              ],
              covariate_feature_time_offset=covariate_delay,
              covariate_feature_window=DEFAULT_WINDOW_SIZE,
              encoder_kwargs={
                  "initial_bias": 0,
                  "location_dependent_bias": True,
                  "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
              }),
      ])
  return result
def get_model_spec_time_varying_with_covariates(covariate_delay):
    """Return model specification with hyperparameters and covariates."""
    result = model_spec_lib.ModelSpec(
        hparams={
            "initial_learning_rate": 0.003061242915556316,
            "momentum": 0.2,
            "decay_steps": 1000,
            "fine_tuning_steps": 100,
            "fine_tuning_decay": 1.0,
            "decay_rate": 1.0,
            "location_dependent_init": False,
            "infected_threshold": 3,
            "restart_threshold": 3000,
            "time_scale_weight": 0.00006243159539906051,
            "train_loss_coefs": [0.01, 0.01, 0.1, 0.3, 0.01, 0.01],
            "valid_loss_coefs": [0.01, 0.01, 0.1, 0.3, 0.01, 0.01],
            "sync_coef": 1.0,
            "reduced_sync_undoc": 1.0,
            "smooth_coef": 0.5,
            "first_dose_vaccine_ratio_per_day_init": 0.0,
            "second_dose_vaccine_ratio_per_day_init": 0.0,
            "average_contact_id_rate_init": -1.9131825459930378,
            "average_contact_iud_rate_init": -1.071866945725303,
            "reinfectable_rate_init": -5.548940468292865,
            "alpha_rate_init": -2.272765554778715,
            "diagnosis_rate_init": -2.095597433974376,
            "recovery_h_rate_init": -1.9032896753850963,
            "recovery_id_rate_init": -1.495660223962899,
            "recovery_iud_rate_init": -1.475605314236803,
            "hospitalization_rate_init": -1.4331763640928012,
            "death_h_rate_init": -1.4886876719378206,
            "death_id_rate_init": -1.8060447968974489,
            "bias_penalty_coef": 0.2835406167308398,
            "r_eff_penalty_coef": 2.0,
            "r_eff_penalty_cutoff": 5.0,
            "acceleration_death_coef": 0.1,
            "acceleration_confirm_coef": 0.1,
            "acceleration_hospital_coef": 0.1,
            "quantile_encoding_window": 7,
            "quantile_smooth_coef": 0.90,
            "quantile_training_iteration_ratio": 1.0,
            "width_coef_train": 1.0,
            "width_coef_valid": 1.0,
            "quantile_cum_viol_coef": 500.0,
            "increment_loss_weight": 0.0,
            "lasso_penalty_coef": 1.0,
            "covariate_training_mixing_coef": 1.0,
            "train_window_range": 2.0,
            "partial_mean_interval": 3,
            "direction_loss_coef": 2000.0,
            "train_crps_weight": 0.75,
            "valid_crps_weight": 0.75
        },
        encoder_specs=[
            model_spec_lib.EncoderSpec(
                encoder_name="first_dose_vaccine_ratio_per_day",
                encoder_type="vaccine",
                vaccine_type="first_dose",
                covariate_feature_specs=[
                    first_dose_vaccine_ratio_per_day,
                    first_dose_vaccine_effectiveness,
                    second_dose_vaccine_effectiveness
                ],
            ),
            model_spec_lib.EncoderSpec(
                encoder_name="second_dose_vaccine_ratio_per_day",
                encoder_type="vaccine",
                vaccine_type="second_dose",
                covariate_feature_specs=[
                    second_dose_vaccine_ratio_per_day,
                    first_dose_vaccine_effectiveness,
                    second_dose_vaccine_effectiveness
                ],
            ),
            model_spec_lib.EncoderSpec(
                encoder_name="average_contact_id_rate",
                encoder_type="gam",
                static_feature_specs=demographics + wellness,
                covariate_feature_specs=[
                    confirmed_preprocessed, death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum, state_of_emergency,
                    reproductive_number
                ] + survey_all + mobility_all,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="average_contact_iud_rate",
                encoder_type="gam",
                static_feature_specs=demographics + wellness,
                covariate_feature_specs=[
                    confirmed_preprocessed, death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum, state_of_emergency,
                    reproductive_number
                ] + survey_all + mobility_all,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(encoder_name="reinfectable_rate",
                                       encoder_type="gam",
                                       static_feature_specs=demographics +
                                       wellness,
                                       covariate_feature_specs=None,
                                       encoder_kwargs={
                                           "initial_bias": 0,
                                           "location_dependent_bias": True,
                                       }),
            model_spec_lib.EncoderSpec(encoder_name="alpha_rate",
                                       encoder_type="gam",
                                       static_feature_specs=[],
                                       covariate_feature_specs=None,
                                       encoder_kwargs={
                                           "initial_bias": 0,
                                           "location_dependent_bias": True,
                                       }),
            model_spec_lib.EncoderSpec(
                encoder_name="diagnosis_rate",
                encoder_type="gam",
                static_feature_specs=healthcare,
                covariate_feature_specs=[
                    confirmed_preprocessed,
                    death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    total_tests,
                    dow_feature,
                ] + survey_all,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="recovery_id_rate",
                encoder_type="gam",
                static_feature_specs=demographics + wellness + healthcare,
                covariate_feature_specs=[
                    confirmed_preprocessed,
                    death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    dow_feature,
                ] + survey_all,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="recovery_iud_rate",
                encoder_type="gam",
                static_feature_specs=demographics + wellness + healthcare,
                covariate_feature_specs=[
                    confirmed_preprocessed,
                    death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                ] + survey_all,
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="recovery_h_rate",
                encoder_type="gam",
                static_feature_specs=demographics + wellness + healthcare,
                covariate_feature_specs=[
                    confirmed_preprocessed,
                    death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    discharged,
                    dow_feature,
                ],
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="hospitalization_rate",
                encoder_type="gam",
                static_feature_specs=demographics + wellness + healthcare,
                covariate_feature_specs=[
                    confirmed_preprocessed,
                    death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    discharged,
                    dow_feature,
                ],
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="death_id_rate",
                encoder_type="gam",
                static_feature_specs=demographics + wellness + healthcare,
                covariate_feature_specs=[
                    confirmed_preprocessed,
                    death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    dow_feature,
                ],
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
            model_spec_lib.EncoderSpec(
                encoder_name="death_h_rate",
                encoder_type="gam",
                static_feature_specs=demographics + wellness + healthcare,
                covariate_feature_specs=[
                    confirmed_preprocessed,
                    death_preprocessed,
                    preprocessed_confirmed_mean_to_sum,
                    preprocessed_death_mean_to_sum,
                    discharged,
                    dow_feature,
                ],
                covariate_feature_time_offset=covariate_delay,
                covariate_feature_window=DEFAULT_WINDOW_SIZE,
                encoder_kwargs={
                    "initial_bias": 0,
                    "location_dependent_bias": True,
                    "use_fixed_covariate_mask": USE_FIXED_COVARIATE_MASK,
                }),
        ])
    return result
Example #5
0
    def get_model_spec(
        self,
        model_type,
        covariate_delay=0,
        **kwargs,
    ):
        # Static feature candidates
        population_density_per_sq_km = model_spec_lib.FeatureSpec(
            name=constants.POPULATION_DENSITY_PER_SQKM, initializer=None)

        # Time-varying feature candidates
        preprocessed_confirmed = model_spec_lib.FeatureSpec(
            name=constants.CONFIRMED_PREPROCESSED,
            initializer=None,
            forecast_method=model_spec_lib.ForecastMethod.NONE)
        preprocessed_death = model_spec_lib.FeatureSpec(
            name=constants.DEATH_PREPROCESSED,
            initializer=None,
            forecast_method=model_spec_lib.ForecastMethod.NONE)
        preprocessed_confirmed_mean_to_sum = model_spec_lib.FeatureSpec(
            name=constants.CONFIRMED_PREPROCESSED_MEAN_TO_SUM_RATIO,
            initializer=None,
            forecast_method=model_spec_lib.ForecastMethod.NONE)
        preprocessed_death_mean_to_sum = model_spec_lib.FeatureSpec(
            name=constants.DEATH_PREPROCESSED_MEAN_TO_SUM_RATIO,
            initializer=None,
            forecast_method=model_spec_lib.ForecastMethod.NONE)
        first_dose_vaccine_ratio_per_day = model_spec_lib.FeatureSpec(
            name=constants.VACCINATED_RATIO_FIRST_DOSE_PER_DAY_PREPROCESSED,
            initializer=None)
        second_dose_vaccine_ratio_per_day = model_spec_lib.FeatureSpec(
            name=constants.VACCINATED_RATIO_SECOND_DOSE_PER_DAY_PREPROCESSED,
            initializer=None)
        first_dose_vaccine_effectiveness = model_spec_lib.FeatureSpec(
            name=constants.VACCINATED_EFFECTIVENESS_FIRST_DOSE,
            initializer=None)
        second_dose_vaccine_effectiveness = model_spec_lib.FeatureSpec(
            name=constants.VACCINATED_EFFECTIVENESS_SECOND_DOSE,
            initializer=None)
        mobility_index = model_spec_lib.FeatureSpec(
            name=constants.MOBILITY_INDEX, initializer=None)

        all_static_covariates = [population_density_per_sq_km]

        base_ts_covariates = [
            preprocessed_confirmed,
            preprocessed_death,
            preprocessed_confirmed_mean_to_sum,
            preprocessed_death_mean_to_sum,
        ]

        default_window_size = 7
        use_fixed_covariate_mask = True

        return model_spec_lib.ModelSpec(
            hparams={
                "initial_learning_rate": 0.003061242915556316,
                "momentum": 0.2,
                "decay_steps": 1000,
                "fine_tuning_steps": 500,
                "fine_tuning_decay": 1.0,
                "decay_rate": 1.0,
                "location_dependent_init": False,
                "infected_threshold": 3,
                "restart_threshold": 300,
                "time_scale_weight": 0.00006243159539906051,
                "train_loss_coefs": [0, 0.001, 0.2, 0.1, 0.01, 0.01],
                "valid_loss_coefs": [0, 0.001, 0.2, 0.1, 0.01, 0.01],
                "sync_coef": 0.3,
                "reduced_sync_undoc": 2.0,
                "smooth_coef": 0.5,
                "first_dose_vaccine_ratio_per_day_init": 0.0,
                "second_dose_vaccine_ratio_per_day_init": 0.0,
                "average_contact_id_rate_init": -1.9131825459930378,
                "average_contact_iud_rate_init": -1.071866945725303,
                "reinfectable_rate_init": -5.548940468292865,
                "alpha_rate_init": -2.272765554778715,
                "diagnosis_rate_init": -2.095597433974376,
                "recovery_id_rate_init": -1.495660223962899,
                "recovery_iud_rate_init": -1.475605314236803,
                "recovery_h_rate_init": -1.9032896753850963,
                "hospitalization_rate_init": -1.4331763640928012,
                "death_id_rate_init": -1.8060447968974489,
                "death_h_rate_init": -1.4886876719378206,
                "bias_penalty_coef": 0.2835406167308398,
                "r_eff_penalty_coef": 2.0,
                "acceleration_death_coef": 0.1,
                "acceleration_confirm_coef": 0.1,
                "acceleration_hospital_coef": 0.1,
                "quantile_encoding_window": 7,
                "quantile_smooth_coef": 0.9,
                "quantile_training_iteration_ratio": 0.5,
                "width_coef_train": 10.0,
                "width_coef_valid": 5.0,
                "quantile_cum_viol_coef": 500.0,
                "increment_loss_weight": 0.0,
                "lasso_penalty_coef": 1.0,
                "covariate_training_mixing_coef": 1.0,
                "train_window_range": 2.0,
                "direction_loss_coef": 5000.0,
                "train_crps_weight": 0.25,
                "valid_crps_weight": 0.25
            },
            encoder_specs=[
                model_spec_lib.EncoderSpec(
                    encoder_name="first_dose_vaccine_ratio_per_day",
                    encoder_type="vaccine",
                    vaccine_type="first_dose",
                    covariate_feature_specs=[
                        first_dose_vaccine_ratio_per_day,
                        first_dose_vaccine_effectiveness,
                        second_dose_vaccine_effectiveness
                    ],
                ),
                model_spec_lib.EncoderSpec(
                    encoder_name="second_dose_vaccine_ratio_per_day",
                    encoder_type="vaccine",
                    vaccine_type="second_dose",
                    covariate_feature_specs=[
                        second_dose_vaccine_ratio_per_day,
                        first_dose_vaccine_effectiveness,
                        second_dose_vaccine_effectiveness
                    ],
                ),
                model_spec_lib.EncoderSpec(
                    encoder_name="average_contact_id_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates +
                    [mobility_index],
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
                model_spec_lib.EncoderSpec(
                    encoder_name="average_contact_iud_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates +
                    [mobility_index],
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
                model_spec_lib.EncoderSpec(
                    encoder_name="reinfectable_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=None,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                    }),
                model_spec_lib.EncoderSpec(encoder_name="alpha_rate",
                                           encoder_type="gam",
                                           static_feature_specs=[],
                                           covariate_feature_specs=None,
                                           encoder_kwargs={
                                               "initial_bias": 0,
                                               "location_dependent_bias": True,
                                           }),
                model_spec_lib.EncoderSpec(
                    encoder_name="diagnosis_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates,
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
                model_spec_lib.EncoderSpec(
                    encoder_name="recovery_id_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates,
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
                model_spec_lib.EncoderSpec(
                    encoder_name="recovery_iud_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates,
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
                model_spec_lib.EncoderSpec(
                    encoder_name="recovery_h_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates,
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
                model_spec_lib.EncoderSpec(
                    encoder_name="hospitalization_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates,
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
                model_spec_lib.EncoderSpec(
                    encoder_name="death_id_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates,
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
                model_spec_lib.EncoderSpec(
                    encoder_name="death_h_rate",
                    encoder_type="gam",
                    static_feature_specs=all_static_covariates,
                    covariate_feature_specs=base_ts_covariates,
                    covariate_feature_window=default_window_size,
                    encoder_kwargs={
                        "initial_bias": 0,
                        "location_dependent_bias": True,
                        "use_fixed_covariate_mask": use_fixed_covariate_mask,
                    }),
            ])