Exemple #1
0
    def test_parameter_sweep_bad_force_initialize(self, model, tmp_path):
        comm, rank, num_procs = _init_mpi()
        tmp_path = _get_rank0_path(comm, tmp_path)

        m = model
        m.fs.slack_penalty = 1000.
        m.fs.slack.setub(0)

        A = m.fs.input['a']
        B = m.fs.input['b']
        sweep_params = {A.name: (A, 0.1, 0.9, 3), B.name: (B, 0.0, 0.5, 3)}

        results_file = os.path.join(tmp_path, 'global_results_recover.csv')
        h5_fname = "output_dict_recover"

        with pytest.raises(ValueError):
            # Call the parameter_sweep function
            parameter_sweep(m,
                            sweep_params,
                            outputs=None,
                            csv_results_file=results_file,
                            h5_results_file=h5_fname,
                            optimize_function=_optimization,
                            reinitialize_before_sweep=True,
                            reinitialize_function=None,
                            reinitialize_kwargs=None,
                            mpi_comm=comm)
Exemple #2
0
def run_parameter_sweep(
    csv_results_file_name=None,
    h5_results_file_name=None,
    seed=None,
    use_LHS=False,
    read_sweep_params_from_file=False,
    sweep_params_fname="mc_sweep_params.yaml",
    read_model_defauls_from_file=False,
    defaults_fname="default_configuration.yaml",
):

    # Set up the solver
    solver = get_solver()

    # Build, set, and initialize the system (these steps will change depending on the underlying model)
    m = build()
    set_operating_conditions(m,
                             water_recovery=0.5,
                             over_pressure=0.3,
                             solver=solver)
    initialize_system(m, solver=solver)

    # Simulate once outside the parameter sweep to ensure everything is appropriately initialized
    solve(m, solver=solver)

    # Check if we need to read in the default model values from a file
    if read_model_defauls_from_file:
        set_defaults_from_yaml(m, defaults_fname)

    # Define the sampling type and ranges for three different variables
    if read_sweep_params_from_file:
        sweep_params = get_sweep_params_from_yaml(m, sweep_params_fname)
    else:
        sweep_params = get_sweep_params(m, use_LHS=use_LHS)

    # Define the outputs to be saved
    outputs = {}
    outputs["EC"] = m.fs.costing.specific_energy_consumption
    outputs["LCOW"] = m.fs.costing.LCOW

    # Run the parameter sweep study using num_samples randomly drawn from the above range
    num_samples = 10

    # Run the parameter sweep
    global_results = parameter_sweep(
        m,
        sweep_params,
        outputs,
        csv_results_file_name=csv_results_file_name,
        h5_results_file_name=h5_results_file_name,
        optimize_function=optimize,
        optimize_kwargs={
            "solver": solver,
            "check_termination": False
        },
        num_samples=num_samples,
        seed=seed,
    )

    return global_results
Exemple #3
0
    def test_parameter_sweep_bad_recover(self, model, tmp_path):
        comm, rank, num_procs = _init_mpi()
        tmp_path = _get_rank0_path(comm, tmp_path)

        m = model
        m.fs.slack_penalty = 1000.
        m.fs.slack.setub(0)

        A = m.fs.input['a']
        B = m.fs.input['b']
        sweep_params = {A.name: (A, 0.1, 0.9, 3), B.name: (B, 0.0, 0.5, 3)}
        results_file = os.path.join(tmp_path, 'global_results_bad_recover.csv')
        h5_fname = "output_dict_bad_recover"

        # Call the parameter_sweep function
        parameter_sweep(m,
                        sweep_params,
                        outputs=None,
                        csv_results_file=results_file,
                        h5_results_file=h5_fname,
                        optimize_function=_optimization,
                        reinitialize_function=_bad_reinitialize,
                        reinitialize_kwargs={'slack_penalty': 10.},
                        mpi_comm=comm)

        # NOTE: rank 0 "owns" tmp_path, so it needs to be
        #       responsible for doing any output file checking
        #       tmp_path can be deleted as soon as this method
        #       returns
        if rank == 0:
            # Check that the global results file is created
            assert os.path.isfile(results_file)

            # Attempt to read in the data
            data = np.genfromtxt(results_file, skip_header=1, delimiter=',')

            # Compare the last row of the imported data to truth
            truth_data = [
                0.9, 0.5, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan
            ]
            assert np.allclose(data[-1], truth_data, equal_nan=True)

        if rank == 0:
            truth_dict = {
                'outputs': {
                    'fs.output[c]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([
                            0.2, 0.2, np.nan, 1., 1., np.nan, np.nan, np.nan,
                            np.nan
                        ])
                    },
                    'fs.output[d]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([
                            0., 0.75, np.nan, 0., 0.75, np.nan, np.nan, np.nan,
                            np.nan
                        ])
                    },
                    'fs.performance': {
                        'value':
                        np.array([
                            0.2, 0.95, np.nan, 1., 1.75, np.nan, np.nan,
                            np.nan, np.nan
                        ])
                    },
                    'fs.slack[ab_slack]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([
                            0., 0., np.nan, 0., 0., np.nan, np.nan, np.nan,
                            np.nan
                        ])
                    },
                    'fs.slack[cd_slack]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([
                            0., 0., np.nan, 0., 0., np.nan, np.nan, np.nan,
                            np.nan
                        ])
                    },
                    'objective': {
                        'value':
                        np.array([
                            0.2, 0.95, np.nan, 1., 1.75, np.nan, np.nan,
                            np.nan, np.nan
                        ])
                    }
                },
                'solve_successful':
                [True, True, False, True, True, False, False, False, False],
                'sweep_params': {
                    'fs.input[a]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([0.1, 0.1, 0.1, 0.5, 0.5, 0.5, 0.9, 0.9, 0.9])
                    },
                    'fs.input[b]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([0., 0.25, 0.5, 0., 0.25, 0.5, 0., 0.25, 0.5])
                    }
                }
            }

            h5_fpath = os.path.join(tmp_path, '{0}.h5'.format(h5_fname))
            read_dict = _read_output_h5(h5_fpath)
            _assert_dictionary_correctness(truth_dict, read_dict)
            _assert_h5_csv_agreement(results_file, read_dict)
Exemple #4
0
    def test_parameter_sweep_optimize(self, model, tmp_path):

        comm, rank, num_procs = _init_mpi()
        tmp_path = _get_rank0_path(comm, tmp_path)

        m = model
        m.fs.slack_penalty = 1000.
        m.fs.slack.setub(0)

        A = m.fs.input['a']
        B = m.fs.input['b']
        sweep_params = {A.name: (A, 0.1, 0.9, 3), B.name: (B, 0.0, 0.5, 3)}
        outputs = {
            'output_c': m.fs.output['c'],
            'output_d': m.fs.output['d'],
            'performance': m.fs.performance,
            'objective': m.objective
        }
        results_file = os.path.join(tmp_path, 'global_results_optimize.csv')
        h5_fname = "output_dict_optimize"

        # Call the parameter_sweep function
        parameter_sweep(m,
                        sweep_params,
                        outputs=outputs,
                        csv_results_file=results_file,
                        h5_results_file=h5_fname,
                        optimize_function=_optimization,
                        optimize_kwargs={'relax_feasibility': True},
                        mpi_comm=comm)

        # NOTE: rank 0 "owns" tmp_path, so it needs to be
        #       responsible for doing any output file checking
        #       tmp_path can be deleted as soon as this method
        #       returns
        if rank == 0:
            # Check that the global results file is created
            assert os.path.isfile(results_file)

            # Attempt to read in the data
            data = np.genfromtxt(results_file, skip_header=1, delimiter=',')
            # Compare the last row of the imported data to truth
            truth_data = [
                0.9, 0.5, 1.0, 1.0, 2.0,
                2.0 - 1000. * ((2. * 0.9 - 1.) + (3. * 0.5 - 1.))
            ]
            assert np.allclose(data[-1], truth_data, equal_nan=True)

        # Check the h5
        if rank == 0:

            truth_dict = {
                'outputs': {
                    'output_c': {
                        'lower bound': 0,
                        'units': 'None',
                        'upper bound': 0,
                        'value':
                        np.array([0.2, 0.2, 0.2, 1., 1., 1., 1., 1., 1.])
                    },
                    'output_d': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([
                            9.98580690e-09, 0.75, 1., 9.99872731e-09, 0.75, 1.,
                            9.99860382e-09, 0.75, 1.
                        ])
                    },
                    'performance': {
                        'value':
                        np.array([0.2, 0.95, 1.2, 1., 1.75, 2., 1., 1.75, 2.])
                    },
                    'objective': {
                        'value':
                        np.array([
                            0.2, 9.50000020e-01, -4.98799990e+02, 1., 1.75,
                            -4.97999990e+02, -7.98999990e+02, -7.98249990e+02,
                            2.0 - 1000. * ((2. * 0.9 - 1.) + (3. * 0.5 - 1.))
                        ])
                    }
                },
                'solve_successful': [True] * 9,
                'sweep_params': {
                    'fs.input[a]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([0.1, 0.1, 0.1, 0.5, 0.5, 0.5, 0.9, 0.9, 0.9])
                    },
                    'fs.input[b]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([0., 0.25, 0.5, 0., 0.25, 0.5, 0., 0.25, 0.5])
                    }
                }
            }

            h5_fpath = os.path.join(tmp_path, '{0}.h5'.format(h5_fname))
            read_dict = _read_output_h5(h5_fpath)
            _assert_dictionary_correctness(truth_dict, read_dict)
            _assert_h5_csv_agreement(results_file, read_dict)
Exemple #5
0
    def test_parameter_sweep(self, model, tmp_path):
        comm, rank, num_procs = _init_mpi()
        tmp_path = _get_rank0_path(comm, tmp_path)

        m = model
        m.fs.slack_penalty = 1000.
        m.fs.slack.setub(0)

        A = m.fs.input['a']
        B = m.fs.input['b']
        sweep_params = {A.name: (A, 0.1, 0.9, 3), B.name: (B, 0.0, 0.5, 3)}
        outputs = {
            'output_c': m.fs.output['c'],
            'output_d': m.fs.output['d'],
            'performance': m.fs.performance
        }
        results_file = os.path.join(tmp_path, 'global_results.csv')
        h5_fname = "output_dict"

        # Call the parameter_sweep function
        parameter_sweep(m,
                        sweep_params,
                        outputs=outputs,
                        csv_results_file=results_file,
                        h5_results_file=h5_fname,
                        optimize_function=_optimization,
                        debugging_data_dir=tmp_path,
                        mpi_comm=comm)

        # NOTE: rank 0 "owns" tmp_path, so it needs to be
        #       responsible for doing any output file checking
        #       tmp_path can be deleted as soon as this method
        #       returns
        if rank == 0:
            # Check that the global results file is created
            assert os.path.isfile(results_file)

            # Check that all local output files have been created
            for k in range(num_procs):
                assert os.path.isfile(
                    os.path.join(tmp_path, f'local_results_{k:03}.csv'))

            # Attempt to read in the data
            data = np.genfromtxt(results_file, skip_header=1, delimiter=',')

            # Compare the last row of the imported data to truth
            truth_data = [0.9, 0.5, np.nan, np.nan, np.nan]
            assert np.allclose(data[-1], truth_data, equal_nan=True)

        # Check for the h5 output
        if rank == 0:
            truth_dict = {
                'outputs': {
                    'output_c': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([
                            0.2, 0.2, np.nan, 1., 1., np.nan, np.nan, np.nan,
                            np.nan
                        ])
                    },
                    'output_d': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([
                            0., 0.75, np.nan, 0., 0.75, np.nan, np.nan, np.nan,
                            np.nan
                        ])
                    },
                    'performance': {
                        'value':
                        np.array([
                            0.2, 0.95, np.nan, 1., 1.75, np.nan, np.nan,
                            np.nan, np.nan
                        ])
                    }
                },
                'solve_successful':
                [True, True, False, True, True, False, False, False, False],
                'sweep_params': {
                    'fs.input[a]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([0.1, 0.1, 0.1, 0.5, 0.5, 0.5, 0.9, 0.9, 0.9])
                    },
                    'fs.input[b]': {
                        'lower bound':
                        0,
                        'units':
                        'None',
                        'upper bound':
                        0,
                        'value':
                        np.array([0., 0.25, 0.5, 0., 0.25, 0.5, 0., 0.25, 0.5])
                    }
                }
            }

            h5_fpath = os.path.join(tmp_path, 'output_dict.h5')
            read_dict = _read_output_h5(h5_fpath)
            _assert_dictionary_correctness(truth_dict, read_dict)
            _assert_h5_csv_agreement(results_file, read_dict)

            # Check if there is a text file created
            import ast

            truth_txt_dict = {
                'outputs': ['output_c', 'output_d', 'performance'],
                'sweep_params': ['fs.input[a]', 'fs.input[b]']
            }

            txt_fpath = os.path.join(tmp_path, '{0}.txt'.format(h5_fname))
            assert os.path.exists(txt_fpath)
            f = open(txt_fpath)
            f_contents = f.read()
            read_txt_dict = ast.literal_eval(f_contents)
            assert read_txt_dict == truth_txt_dict
Exemple #6
0
def run_analysis(case_num, nx, interpolate_nan_outputs=True):
    m, _ = metab.main()

    outputs, optimize_kwargs, opt_function = set_up_sensitivity(m)

    sweep_params = {}
    if case_num == 1:
        # bead cost
        sweep_params["bead_cost"] = LinearSample(
            m.fs.costing.metab.bead_cost["hydrogen"], 1, 50, nx)

    elif case_num == 2:
        # bead replacement rate
        # baseline corresponds to replacement rate of 0.3 years; sensitivity on replacement_factor corresponding
        # to 0.3 yr to 5 yr
        import numpy as np
        from pyomo.environ import value

        upper_replacement_rate = 5  # replace every x years
        replacement_intervals = np.arange(
            upper_replacement_rate,
            value(m.fs.costing.plant_lifetime),
            upper_replacement_rate,
        )
        rep_factor_upper_lim = sum(1 / (1 + value(m.fs.costing.wacc))**x
                                   for x in replacement_intervals) * value(
                                       m.fs.costing.capital_recovery_factor)

        sweep_params["bead_cost"] = LinearSample(
            m.fs.costing.metab.bead_replacement_factor["hydrogen"],
            3.376,
            rep_factor_upper_lim,
            nx,
        )

    elif case_num == 3:
        # Hydrogen METAB HRT
        sweep_params["hydrogen_hrt"] = LinearSample(
            m.fs.metab_hydrogen.hydraulic_retention_time, 0.75, 24, nx)

    elif case_num == 4:
        # Methane METAB HRT
        sweep_params["methane_hrt"] = LinearSample(
            m.fs.metab_methane.hydraulic_retention_time, 47.25, 360, nx)

    elif case_num == 5:
        # Hydrogen Conversion Rate: sweep from 0.06 to 0.6 L H2/g-COD removed
        sweep_params["hydrogen_conversion_rate"] = LinearSample(
            m.fs.metab_hydrogen.generation_ratio["cod_to_hydrogen",
                                                 "hydrogen"],
            5.03e-3,
            5e-2,
            nx,
        )

    elif case_num == 6:
        # Methane Conversion Rate: sweep from ~ 0.1 to 0.3 L H2/g-COD removed
        sweep_params["methane_conversion_rate"] = LinearSample(
            m.fs.metab_methane.generation_ratio["cod_to_methane", "methane"],
            6.68e-2,
            2e-1,
            nx,
        )
    elif case_num == 7:
        # Recovered resource selling price:
        sweep_params["Hydrogen selling price"] = LinearSample(
            m.fs.costing.hydrogen_product_cost, -2, -10, nx)
        sweep_params["Methane selling price"] = LinearSample(
            m.fs.costing.methane_product_cost, -0.305, -1, nx)

    else:
        raise ValueError("case_num = %d not recognized." % (case_num))

    output_filename = "sensitivity_" + str(case_num) + ".csv"

    global_results = parameter_sweep(
        m,
        sweep_params,
        outputs,
        csv_results_file_name=output_filename,
        optimize_function=opt_function,
        optimize_kwargs=optimize_kwargs,
        # debugging_data_dir=os.path.split(output_filename)[0] + '/local',
        interpolate_nan_outputs=interpolate_nan_outputs,
    )

    return global_results, sweep_params
Exemple #7
0
def run_analysis(case_num,
                 nx,
                 RO_type,
                 output_directory=None,
                 interp_nan_outputs=True):

    desal_kwargs = {
        "has_desal_feed": False,
        "is_twostage": True,
        "has_ERD": True,
        "RO_type": RO_type,
        "RO_base": "TDS",
        "RO_level": "detailed",
    }

    sweep_params = {}
    outputs = {}
    optimize_kwargs = {"check_termination": False}  # None
    base_path = get_output_directory(output_directory)

    if case_num == 1:
        # ================================================================
        # Single Stage RO
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_single_stage as fs_single_stage

        desal_kwargs["is_twostage"] = False

        m = fs_single_stage.optimize_flowsheet(**desal_kwargs)

        sweep_params["System Recovery"] = LinearSample(
            m.fs.system_recovery_target, 0.3, 0.95, nx)
        # sweep_params['RO Recovery'] = LinearSample(m.fs.RO_recovery, 0.3, 0.95, nx)

        outputs["LCOW"] = m.fs.costing.LCOW
        outputs["Saturation Index"] = m.fs.desal_saturation.saturation_index
        outputs["Pump Pressure"] = m.fs.pump_RO.control_volume.properties_out[
            0].pressure
        outputs["RO Recovery"] = m.fs.RO_recovery
        outputs[
            "Annual Water Production"] = m.fs.costing.annual_water_production
        outputs = append_costing_outputs(m, outputs, ["RO", "pump_RO", "ERD"])

        output_filename = (
            base_path +
            f"output/fs_single_stage/results_{case_num}_{desal_kwargs['RO_type']}RO.csv"
        )

        opt_function = fs_single_stage.optimize

    elif case_num == 2:
        # ================================================================
        # Two Stage RO
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_two_stage as fs_two_stage

        m = fs_two_stage.optimize_flowsheet(**desal_kwargs)

        sweep_params["System Recovery"] = LinearSample(
            m.fs.system_recovery_target, 0.3, 0.95, nx)
        # sweep_params['RO Recovery'] = LinearSample(m.fs.RO_recovery, 0.3, 0.95, nx)
        sweep_params["Max Pressure"] = LinearSample(
            m.fs.max_allowable_pressure, 300e5, 75e5, nx)

        outputs["LCOW"] = m.fs.costing.LCOW
        outputs["Saturation Index"] = m.fs.desal_saturation.saturation_index
        outputs[
            "RO-1 Pump Pressure"] = m.fs.pump_RO.control_volume.properties_out[
                0].pressure
        outputs[
            "RO-2 Pump Pressure"] = m.fs.pump_RO2.control_volume.properties_out[
                0].pressure
        outputs["RO Recovery"] = m.fs.RO_recovery
        outputs[
            "Annual Water Production"] = m.fs.costing.annual_water_production
        outputs = append_costing_outputs(
            m, outputs, ["RO", "pump_RO", "RO2", "pump_RO2", "ERD"])

        output_filename = (
            base_path +
            f"output/fs_two_stage/results_{case_num}_{desal_kwargs['RO_type']}RO.csv"
        )

        opt_function = fs_two_stage.optimize

    elif case_num == 3:
        # ================================================================
        # NF No Bypass
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_NF_no_bypass as fs_NF_no_bypass

        m = fs_NF_no_bypass.solve_flowsheet()

        sweep_params["NF Recovery"] = LinearSample(
            m.fs.NF.recovery_mass_phase_comp[0, "Liq", "H2O"], 0.3, 0.95, nx)

        outputs["LCOW"] = m.fs.costing.LCOW
        outputs["Saturation Index"] = m.fs.pretrt_saturation.saturation_index

        output_filename = base_path + f"output/fs_NF_no_bypass/results_{case_num}.csv"

        opt_function = fs_NF_no_bypass.simulate
        # Need to unfix NF area to simulate with fixed NF recovery
        optimize_kwargs = {"unfix_nf_area": True, "check_termination": False}

    elif case_num == 4:
        # ================================================================
        # NF
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_NF as fs_NF

        m = fs_NF.optimize_flowsheet()

        sweep_params["NF Split Fraction"] = LinearSample(
            m.fs.splitter.split_fraction[0, "pretreatment"], 0.01, 0.99, nx)

        outputs["LCOW"] = m.fs.costing.LCOW
        outputs["Ca Removal"] = m.fs.removal_Ca
        outputs["Max Concentration Constraint"] = m.fs.eq_max_conc_NF.body

        output_filename = (
            base_path +
            f"output/fs_NF/results_{case_num}_{desal_kwargs['RO_type']}RO.csv")

        opt_function = fs_NF.optimize

    elif case_num == 5:  # pragma: no cover
        # ================================================================
        # NF Two Stage
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_NF_two_stage as fs_NF_two_stage

        m = fs_NF_two_stage.optimize_flowsheet(**desal_kwargs)

        sweep_params["NF Split Fraction"] = LinearSample(
            m.fs.splitter.split_fraction[0, "pretreatment"], 0.25, 0.99, 4)
        # sweep_params['RO Recovery'] = LinearSample(m.fs.RO_recovery, 0.3, 0.95, nx)
        sweep_params["System Recovery"] = LinearSample(
            m.fs.system_recovery_target, 0.5, 0.95, nx)

        outputs["LCOW"] = m.fs.costing.LCOW
        outputs["RO Recovery"] = m.fs.RO_recovery

        output_filename = (
            base_path +
            f"output/fs_NF_two_stage/results_{case_num}_{desal_kwargs['RO_type']}RO.csv"
        )

        opt_function = fs_NF_two_stage.optimize

    elif case_num == 6:
        # ================================================================
        # NF Two Stage
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_NF_two_stage as fs_NF_two_stage

        m = fs_NF_two_stage.optimize_flowsheet(**desal_kwargs)

        sweep_params["System Recovery"] = LinearSample(
            m.fs.system_recovery_target, 0.5, 0.95, nx)

        outputs["LCOW"] = m.fs.costing.LCOW
        outputs["NF Split Fraction"] = m.fs.splitter.split_fraction[
            0, "pretreatment"]

        output_filename = (
            base_path +
            f"output/fs_NF_two_stage/results_{case_num}_{desal_kwargs['RO_type']}RO.csv"
        )

        opt_function = fs_NF_two_stage.optimize

    elif case_num == 7:  # pragma: no cover
        # ================================================================
        # NF Two Stage
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_NF_two_stage as fs_NF_two_stage

        m = fs_NF_two_stage.optimize_flowsheet(**desal_kwargs)

        sweep_params["NF Split Fraction"] = LinearSample(
            m.fs.splitter.split_fraction[0, "pretreatment"], 0.25, 0.99, 4)
        sweep_params["System Recovery"] = LinearSample(
            m.fs.system_recovery_target, 0.5, 0.95, nx)

        outputs["LCOW"] = m.fs.costing.LCOW

        output_filename = (
            base_path +
            f"output/fs_NF_two_stage/results_{case_num}_{desal_kwargs['RO_type']}RO.csv"
        )

        opt_function = fs_NF_two_stage.optimize

    elif case_num == 8:
        # ================================================================
        # Softening
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_softening as fs_softening

        m = fs_softening.solve_flowsheet()

        sweep_params["Lime Dosing"] = LinearSample(
            m.fs.stoich_softening_mixer_unit.lime_stream_state[0].flow_mol,
            0.0,
            0.128,
            nx,
        )

        outputs["LCOW"] = m.fs.costing.LCOW
        outputs["Ca Removal"] = m.fs.removal_Ca
        outputs["Mg Removal"] = m.fs.removal_Mg
        outputs[
            "Annual Water Production"] = m.fs.costing.annual_water_production
        outputs["capital_cost_total"] = m.fs.costing.total_capital_cost
        outputs["operating_cost_total"] = m.fs.costing.total_operating_cost

        output_filename = base_path + f"output/fs_softening/results_{case_num}.csv"

        opt_function = fs_softening.simulate

    elif case_num == 9:
        # ================================================================
        # Softening Two Stage
        # ================================================================
        import watertap.examples.flowsheets.full_treatment_train.analysis.flowsheet_softening_two_stage as fs_softening_two_stage

        m = fs_softening_two_stage.optimize_flowsheet(**desal_kwargs)

        sweep_params["System Recovery"] = LinearSample(
            m.fs.system_recovery_target, 0.5, 0.95, nx)

        outputs["LCOW"] = m.fs.costing.LCOW
        outputs[
            "Lime Dosing"] = m.fs.stoich_softening_mixer_unit.lime_stream_state[
                0].flow_mol

        output_filename = (
            base_path +
            f"output/fs_softening_two_stage/results_{case_num}_{desal_kwargs['RO_type']}RO.csv"
        )

        opt_function = fs_softening_two_stage.optimize

    else:
        raise ValueError("case_num = %d not recognized." % (case_num))

    global_results = parameter_sweep(
        m,
        sweep_params,
        outputs,
        csv_results_file_name=output_filename,
        optimize_function=opt_function,
        optimize_kwargs=optimize_kwargs,
        debugging_data_dir=os.path.split(output_filename)[0] + "/local",
        interpolate_nan_outputs=interp_nan_outputs,
    )

    return global_results, sweep_params