From c92bdb85a6e8b9cb640dafc163eee179e36baa11 Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 11 Aug 2025 07:52:15 +0200 Subject: [PATCH 01/41] Extrapolating and saving real world data, initialize model in data generation based on real world ratios --- .../memilio/surrogatemodel/GNN/GNN_utils.py | 160 +++++++ .../surrogatemodel/GNN/data_generation.py | 400 ++++++++++++++++++ .../GNN/extrapolate_data_gnn.py | 223 ++++++++++ 3 files changed, 783 insertions(+) create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py new file mode 100644 index 0000000000..53c30e79db --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -0,0 +1,160 @@ +import numpy as np +import pandas as pd +import os +from sklearn.preprocessing import FunctionTransformer + +from memilio.epidata import transformMobilityData as tmd +from memilio.epidata import getDataIntoPandasDataFrame as gd +from memilio.simulation.osecir import (ModelGraph, set_edges) +from memilio.epidata import modifyDataframeSeries as mdfs + + +def remove_confirmed_compartments(dataset_entries, num_groups): + """ The compartments which contain confirmed cases are not needed and are + therefore omitted by summarizing the confirmed compartment with the + original compartment. + + :param dataset_entries: Array that contains the compartmental data with + confirmed compartments. + :param num_groups: Number of age groups. + :returns: Array that contains the compartmental data without confirmed compartments. + """ + + new_dataset_entries = [] + for i in dataset_entries: + dataset_entries_reshaped = i.reshape( + [num_groups, int(np.asarray(dataset_entries).shape[1]/num_groups)] + ) + sum_inf_no_symp = np.sum(dataset_entries_reshaped[:, [2, 3]], axis=1) + sum_inf_symp = np.sum(dataset_entries_reshaped[:, [4, 5]], axis=1) + dataset_entries_reshaped[:, 2] = sum_inf_no_symp + dataset_entries_reshaped[:, 4] = sum_inf_symp + new_dataset_entries.append( + np.delete(dataset_entries_reshaped, [3, 5], axis=1).flatten() + ) + return new_dataset_entries + + +def getBaselineMatrix(): + """ loads the baselinematrix + """ + + baseline_contact_matrix0 = os.path.join( + "./data/Germany/contacts/baseline_home.txt") + baseline_contact_matrix1 = os.path.join( + "./data/Germany/contacts/baseline_school_pf_eig.txt") + baseline_contact_matrix2 = os.path.join( + "./data/Germany/contacts/baseline_work.txt") + baseline_contact_matrix3 = os.path.join( + "./data/Germany/contacts/baseline_other.txt") + + baseline = np.loadtxt(baseline_contact_matrix0) \ + + np.loadtxt(baseline_contact_matrix1) + \ + np.loadtxt(baseline_contact_matrix2) + \ + np.loadtxt(baseline_contact_matrix3) + + return baseline + + +def make_graph(directory, num_regions, countykey_list, models): + """ Generating graph with one node per region. + + Each node contains the county ID and a osecir model for the county. The edges + contain the mobility information between different counties/nodes. + + :param directory: Directory with mobility data. + :param num_regions: Number (int) of counties that should be added to the + graph-ODE model. Equals 400 for whole Germany. + :param countykey_list: List of keys/IDs for each county. + :models models: List of osecir Model with one model per county. + :returns: Graph-ODE model. + """ + graph = ModelGraph() + for i in range(num_regions): + graph.add_node(int(countykey_list[i]), models[i]) + + num_locations = 4 + + set_edges(os.path.abspath(os.path.join(directory, os.pardir)), + graph, num_locations) + return graph + + +def transform_mobility_directory(): + """ Transforms the mobility data by merging Eisenach and Wartburgkreis + """ + # get mobility data directory + arg_dict = gd.cli("commuter_official") + + directory = arg_dict['out_folder'].split('/pydata')[0] + directory = os.path.join(directory, 'Germany/mobility/') + + # Merge Eisenach and Wartbugkreis in Input Data + tmd.updateMobility2022(directory, mobility_file='twitter_scaled_1252') + tmd.updateMobility2022( + directory, mobility_file='commuter_mobility_2022') + return directory + + +def get_population(): + """ + Loading the population data for the different counties and the different age groups. + """ + df_population = pd.read_json( + 'data/Germany/pydata/county_population.json') + age_groups = ['0-4', '5-14', '15-34', '35-59', '60-79', '80-130'] + + df_population_agegroups = pd.DataFrame( + columns=[df_population.columns[0]] + age_groups) + for region_id in df_population.iloc[:, 0]: + df_population_agegroups.loc[len(df_population_agegroups.index), :] = [int(region_id)] + list( + mdfs.fit_age_group_intervals(df_population[df_population.iloc[:, 0] == int(region_id)].iloc[:, 2:], age_groups)) + + population = df_population_agegroups.values.tolist() + return population + + +def scale_data(data): + """ Apply a logarithmic transformation on the data. + + :param data: dictionary, containing entries "inputs" and "labels" + :returns scaled_inputs: Transformed input data + :returns scaled_labels: Transformed output data + """ + + num_groups = int(np.asarray(data['inputs']).shape[2] / 8) + transformer = FunctionTransformer(np.log1p, validate=True) + + # Scale inputs + inputs = np.asarray( + data['inputs']).transpose(2, 0, 1, 3).reshape(num_groups * 8, -1) + scaled_inputs = transformer.transform(inputs) + original_shape_input = np.asarray(data['inputs']).shape + + # Reverse the reshape + reshaped_back = scaled_inputs.reshape(original_shape_input[2], + original_shape_input[0], + original_shape_input[1], + original_shape_input[3]) + + # Reverse the transpose + original_inputs = reshaped_back.transpose(1, 2, 0, 3) + scaled_inputs = original_inputs.transpose(0, 3, 1, 2) + + # Scale labels + labels = np.asarray( + data['labels']).transpose(2, 0, 1, 3).reshape(num_groups * 8, -1) + scaled_labels = transformer.transform(labels) + original_shape_labels = np.asarray(data['labels']).shape + + # Reverse the reshape + reshaped_back = scaled_labels.reshape(original_shape_labels[2], + original_shape_labels[0], + original_shape_labels[1], + original_shape_labels[3]) + + # Reverse the transpose + original_labels = reshaped_back.transpose(1, 2, 0, 3) + scaled_labels = original_labels.transpose(0, 3, 1, 2) + + return scaled_inputs, scaled_labels diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py new file mode 100644 index 0000000000..ee6a97edb3 --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -0,0 +1,400 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Agatha Schmidt, Henrik Zunker +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + +import copy +import os +import pickle +import random +import time +import memilio.simulation as mio +import memilio.simulation.osecir as osecir +import numpy as np + +from progress.bar import Bar + +from memilio.simulation import (AgeGroup, LogLevel, set_log_level, Damping) +from memilio.simulation.osecir import (Index_InfectionState, interpolate_simulation_result, ParameterStudy, + InfectionState, Model, interpolate_simulation_result) + +from memilio.surrogatemodel.GNN.GNN_utils import (transform_mobility_directory, + make_graph, scale_data, getBaselineMatrix, remove_confirmed_compartments) +import memilio.surrogatemodel.utils.dampings as dampings + +from enum import Enum + + +# Enumerate the different locations +class Location(Enum): + Home = 0 + School = 1 + Work = 2 + Other = 3 + + +start_date = mio.Date(2019, 1, 1) +end_date = mio.Date(2021, 12, 31) + + +def set_covid_parameters(model, num_groups=6): + """Setting COVID-parameters for the different age groups. + + :param model: memilio model, whose parameters should be clarified + :param num_groups: Number of age groups + """ + + # age specific parameters + TransmissionProbabilityOnContact = [0.03, 0.06, 0.06, 0.06, 0.09, 0.175] + RecoveredPerInfectedNoSymptoms = [0.25, 0.25, 0.2, 0.2, 0.2, 0.2] + SeverePerInfectedSymptoms = [0.0075, 0.0075, 0.019, 0.0615, 0.165, 0.225] + CriticalPerSevere = [0.075, 0.075, 0.075, 0.15, 0.3, 0.4] + DeathsPerCritical = [0.05, 0.05, 0.14, 0.14, 0.4, 0.6] + + TimeInfectedNoSymptoms = [2.74, 2.74, 2.565, 2.565, 2.565, 2.565] + TimeInfectedSymptoms = [7.02625, 7.02625, + 7.0665, 6.9385, 6.835, 6.775] + TimeInfectedSevere = [5, 5, 5.925, 7.55, 8.5, 11] + TimeInfectedCritical = [6.95, 6.95, 6.86, 17.36, 17.1, 11.6] + + for i, rho, muCR, muHI, muUH, muDU, tc, ti, th, tu in zip(range(num_groups), + TransmissionProbabilityOnContact, RecoveredPerInfectedNoSymptoms, + SeverePerInfectedSymptoms, CriticalPerSevere, DeathsPerCritical, + TimeInfectedNoSymptoms, TimeInfectedSymptoms, + TimeInfectedSevere, TimeInfectedCritical): + # Compartment transition duration + model.parameters.TimeExposed[AgeGroup(i)] = 3.335 + model.parameters.TimeInfectedNoSymptoms[AgeGroup(i)] = tc + model.parameters.TimeInfectedSymptoms[AgeGroup(i)] = ti + model.parameters.TimeInfectedSevere[AgeGroup(i)] = th + model.parameters.TimeInfectedCritical[AgeGroup(i)] = tu + + # Compartment transition propabilities + model.parameters.RelativeTransmissionNoSymptoms[AgeGroup(i)] = 1 + model.parameters.TransmissionProbabilityOnContact[AgeGroup(i)] = rho + model.parameters.RecoveredPerInfectedNoSymptoms[AgeGroup(i)] = muCR + model.parameters.RiskOfInfectionFromSymptomatic[AgeGroup(i)] = 0.25 + model.parameters.SeverePerInfectedSymptoms[AgeGroup(i)] = muHI + model.parameters.CriticalPerSevere[AgeGroup(i)] = muUH + model.parameters.DeathsPerCritical[AgeGroup(i)] = muDU + # twice the value of RiskOfInfectionFromSymptomatic + model.parameters.MaxRiskOfInfectionFromSymptomatic[AgeGroup(i)] = 0.5 + # StartDay is the n-th day of the year + model.parameters.StartDay = start_date.day_in_year + + +def set_contact_matrices(model, data_dir, num_groups=6): + """Setting the contact matrices for a model + + :param model: memilio ODE-model, whose contact matrices should be modified. + :param data_dir: directory, where the contact data is stored (should contain folder "contacts") + :param num_groups: Number of age groups considered + + """ + + contact_matrices = mio.ContactMatrixGroup( + len(list(Location)), num_groups) + locations = ["home", "school_pf_eig", "work", "other"] + + # Loading contact matrices for each location from .txt file + for i, location in enumerate(locations): + baseline_file = os.path.join( + data_dir, "Germany", "contacts", "baseline_" + location + ".txt") + + contact_matrices[i] = mio.ContactMatrix( + mio.read_mobility_plain(baseline_file), + ) + model.parameters.ContactPatterns.cont_freq_mat = contact_matrices + + +def get_graph(num_groups, data_dir, mobility_directory): + """ Generate the associated graph given the mobility data. + + :param num_groups: Number of age groups + :param data_dir: Directory, where the contact data is stored (should contain folder "contacts") + :param mobility_directory: Directory containing the mobility data + """ + # Generating and Initializing the model + model = Model(num_groups) + set_covid_parameters(model) + set_contact_matrices(model, data_dir) + + # Generating the graph + graph = osecir.ModelGraph() + + # Setting the parameters + scaling_factor_infected = [2.5, 2.5, 2.5, 2.5, 2.5, 2.5] + scaling_factor_icu = 1.0 + tnt_capacity_factor = 7.5 / 100000. + + # Path containing the population data + data_dir_Germany = os.path.join(data_dir, "Germany") + pydata_dir = os.path.join(data_dir_Germany, "pydata") + + path_population_data = os.path.join(pydata_dir, + "county_current_population.json") + + # Setting node information based on model parameters + mio.osecir.set_nodes( + model.parameters, + mio.Date(start_date.year, + start_date.month, start_date.day), + mio.Date(end_date.year, + end_date.month, end_date.day), pydata_dir, + path_population_data, True, graph, scaling_factor_infected, + scaling_factor_icu, tnt_capacity_factor, 0, False) + + # Setting edge information based on the mobility data + mio.osecir.set_edges( + mobility_directory, graph, len(Location)) + + return graph + + +def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_groups=6, start_date=mio.Date(2020, 6, 1)): + """ Uses an ODE SECIR model allowing for asymptomatic infection with 6 + different age groups. The model is not stratified by region. + Virus-specific parameters are fixed and initial number of person + in the particular infection states are chosen randomly from defined ranges. + + :param days: Number of days simulated within a single run. + :param damping_days: Days, where a damping is applied + :param damping_factors: damping factors associated to the damping days + :param graph: Graph initialized for the start_date with the population data which + is sampled during the run. + :param num_groups: Number of age groups considered in the simulation + :param start_date: Date, when the simulation starts + :returns: List containing the populations in each compartment used to initialize + the run. + """ + if len(damping_days) != len(damping_factors): + raise ValueError("Length of damping_days and damping_factors differ!") + + min_date = mio.Date(2020, 6, 1) + min_date_num = min_date.day_in_year + start_date_num = start_date.day_in_year - min_date_num + + # Load the ground truth data + with open("/localdata1/hege_mn/memilio/data/Germany/pydata/ground_truth_all_nodes.pickle", 'rb') as f: + ground_truth_all_nodes_np = pickle.load(f) + + # Initialize model for each node, using the population data and sampling the number of + # individuals in the different compartments + for node_indx in range(graph.num_nodes): + model = graph.get_node(node_indx).property + data = ground_truth_all_nodes_np[node_indx][start_date_num] + + # Set parameters + # TODO: Put This in the draw_sample function in the ParameterStudy + # Iterate over the different age groups + for i in range(num_groups): + age_group = AgeGroup(i) + pop_age_group = model.populations.get_group_total_AgeGroup( + age_group) + + valid_configuration = False + while valid_configuration is False: + comp_data = data[:, i] + ratios = np.random.uniform(0.8, 1.2, size=8) + modified_data = comp_data * ratios + if np.sum(modified_data[1:]) < pop_age_group: + valid_configuration = True + + # Set the populations for the different compartments + model.populations[age_group, Index_InfectionState( + InfectionState.Exposed)] = modified_data[1] + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedNoSymptoms)] = modified_data[2] + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedSymptoms)] = modified_data[3] + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedSevere)] = modified_data[4] + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedCritical)] = modified_data[5] + model.populations[age_group, Index_InfectionState( + InfectionState.Recovered)] = modified_data[6] + model.populations[age_group, Index_InfectionState( + InfectionState.Dead)] = modified_data[7] + model.populations.set_difference_from_group_total_AgeGroup(( + age_group, InfectionState.Susceptible), pop_age_group) + + # Introduce the damping information, in general dampings can be local, but till now the code just allows global dampings + damped_matrices = [] + damping_coefficients = [] + + for i in np.arange(len(damping_days)): + day = damping_days[i] + factor = damping_factors[i] + + damping = np.ones((num_groups, num_groups) + ) * np.float16(factor) + model.parameters.ContactPatterns.cont_freq_mat.add_damping(Damping( + coeffs=(damping), t=day, level=0, type=0)) +# damping matrices and damping_coefficients werden überschrieben -> am Ende nur von letztem Knoten übergeben ??? + damped_matrices.append(model.parameters.ContactPatterns.cont_freq_mat.get_matrix_at( + day+1)) + damping_coefficients.append(damping) + + # Apply mathematical constraints to parameters + model.apply_constraints() + + # set model to graph + graph.get_node(node_indx).property.populations = model.populations + + # Start simulation + study = ParameterStudy(graph, 0, days, dt=0.5, num_runs=1) + start_time = time.perf_counter() + study.run() + runtime = time.perf_counter() - start_time + + graph_run = study.run()[0] + results = interpolate_simulation_result(graph_run) + + for result_indx in range(len(results)): + results[result_indx] = remove_confirmed_compartments( + np.asarray(results[result_indx]), num_groups) + + dataset_entry = copy.deepcopy(results) + + return dataset_entry, damped_matrices, damping_coefficients, runtime + + +def generate_data( + num_runs, data_dir, path, input_width, label_width, save_data=True, + transform=True, damping_method="classic", max_number_damping=3): + """ Generate dataset by calling run_secir_simulation (num_runs)-often + + :param num_runs: Number of times, the function run_secir_simulation is called. + :param data_dir: Directory with all data needed to initialize the models. + :param path: Path, where the datasets are stored. + :param input_width: number of time steps used for model input. + :param label_width: number of time steps (days) used as model output/label. + :param save_data: Option to deactivate the save of the dataset. Per default True. + :param damping_method: String specifying the damping method, that should be used. Possible values "classic", "active", "random". + :param max_number_damping: Maximal number of possible dampings. + :returns: Dictionary containing inputs, labels, contact_matrix, damping_days and damping_factors + """ + set_log_level(mio.LogLevel.Error) + days = label_width + input_width - 1 + + # Preparing output dictionary + data = { + "inputs": [], + "labels": [], + "contact_matrix": [], + "damping_days": [], + "damping_factors": [] + } + + # Setting basic parameter + num_groups = 6 + mobility_dir = "/localdata1/hege_mn/memilio/data/Germany/mobility/commuter_mobility_2022.txt" + graph = get_graph(num_groups, data_dir, mobility_dir) + start_dates = [ + mio.Date(2020, 6, 1), + mio.Date(2020, 7, 1), + mio.Date(2020, 8, 1), + mio.Date(2020, 9, 1), + mio.Date(2020, 10, 1), + mio.Date(2020, 11, 1), + mio.Date(2020, 12, 1) + ] + + # show progess in terminal for longer runs + # Due to the random structure, theres currently no need to shuffle the data + bar = Bar('Number of Runs done', max=num_runs) + + times = [] + for i in range(0, num_runs): + start_date = start_dates[i % len(start_dates)] + # Generate random damping days and damping factors + damping_days, damping_factors = dampings.generate_dampings( + days, max_number_damping, method=damping_method, min_distance=2, + min_damping_day=2) + # Run simulation + data_run, damped_matrices, damping_coefficients, t_run = run_secir_groups_simulation( + days, damping_days, damping_factors, graph, num_groups, start_date) + + times.append(t_run) + + inputs = np.asarray(data_run).transpose(1, 0, 2)[: input_width] + labels = np.asarray(data_run).transpose(1, 0, 2)[input_width:] + + data["inputs"].append(inputs) + data["labels"].append(labels) + data["contact_matrix"].append(np.array(damped_matrices)) + data["damping_factors"].append(damping_coefficients) + data["damping_days"].append(damping_days) + + bar.next() + + bar.finish() + + print( + f"For Days = {days}, AVG runtime: {np.mean(times)}s, Median runtime: {np.median(times)}s") + + if save_data: + if transform: + inputs, labels = scale_data(data) + + all_data = {"inputs": scaled_inputs, + "labels": scaled_labels, + "damping_day": data["damping_days"], + "contact_matrix": data["contact_matrix"], + "damping_coeff": data["damping_factors"] + } + + # check if data directory exists. If necessary create it. + if not os.path.isdir(path): + os.mkdir(path) + + # generate the filename + if num_runs < 1000: + filename = 'GNN_data_%ddays_%ddampings_' % ( + label_width, max_number_damping) + damping_method+'%d.pickle' % (num_runs) + else: + filename = 'GNN_data_%ddays_%ddampings_' % ( + label_width, max_number_damping) + damping_method+'%dk.pickle' % (num_runs//1000) + + # save dict to pickle file + with open(os.path.join(path, filename), 'wb') as f: + pickle.dump(all_data, f) + + return data + + +if __name__ == "__main__": + + path = os.getcwd() + path_output = os.path.join(os.getcwd(), 'saves') + # data_dir = os.path.join(os.getcwd(), 'data') + data_dir = "/localdata1/hege_mn/memilio/data" + input_width = 5 + number_of_dampings = 3 + num_runs = 30 + label_width_list = [30] + + random.seed(10) + for label_width in label_width_list: + generate_data(num_runs=num_runs, + data_dir=data_dir, + path=path_output, + input_width=input_width, + label_width=label_width, + save_data=False, + max_number_damping=number_of_dampings) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py new file mode 100644 index 0000000000..3a6f84b9b0 --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py @@ -0,0 +1,223 @@ +import os +import numpy as np +import pickle +import h5py +import pandas as pd +import matplotlib.pyplot as plt +import memilio.simulation as mio +import tensorflow as tf + +from enum import Enum +from memilio.simulation import (AgeGroup, LogLevel, set_log_level, Damping) +from memilio.simulation.osecir import ( + InfectionState, Model, ModelGraph, set_nodes) +from memilio.surrogatemodel.GNN.GNN_utils import remove_confirmed_compartments + + +class Location(Enum): + Home = 0 + School = 1 + Work = 2 + Other = 3 + + +def set_covid_parameters(model, start_date, num_groups=6): + """Setting COVID-parameters for the different age groups. + + :param model: memilio model, whose parameters should be clarified + :param num_groups: Number of age groups + """ + + # age specific parameters + TransmissionProbabilityOnContact = [0.03, 0.06, 0.06, 0.06, 0.09, 0.175] + RecoveredPerInfectedNoSymptoms = [0.25, 0.25, 0.2, 0.2, 0.2, 0.2] + SeverePerInfectedSymptoms = [0.0075, 0.0075, 0.019, 0.0615, 0.165, 0.225] + CriticalPerSevere = [0.075, 0.075, 0.075, 0.15, 0.3, 0.4] + DeathsPerCritical = [0.05, 0.05, 0.14, 0.14, 0.4, 0.6] + + TimeInfectedNoSymptoms = [2.74, 2.74, 2.565, 2.565, 2.565, 2.565] + TimeInfectedSymptoms = [7.02625, 7.02625, + 7.0665, 6.9385, 6.835, 6.775] + TimeInfectedSevere = [5, 5, 5.925, 7.55, 8.5, 11] + TimeInfectedCritical = [6.95, 6.95, 6.86, 17.36, 17.1, 11.6] + + for i, rho, muCR, muHI, muUH, muDU, tc, ti, th, tu in zip(range(num_groups), + TransmissionProbabilityOnContact, RecoveredPerInfectedNoSymptoms, + SeverePerInfectedSymptoms, CriticalPerSevere, DeathsPerCritical, + TimeInfectedNoSymptoms, TimeInfectedSymptoms, + TimeInfectedSevere, TimeInfectedCritical): + # Compartment transition duration + model.parameters.TimeExposed[AgeGroup(i)] = 3.335 + model.parameters.TimeInfectedNoSymptoms[AgeGroup(i)] = tc + model.parameters.TimeInfectedSymptoms[AgeGroup(i)] = ti + model.parameters.TimeInfectedSevere[AgeGroup(i)] = th + model.parameters.TimeInfectedCritical[AgeGroup(i)] = tu + + # Compartment transition propabilities + model.parameters.RelativeTransmissionNoSymptoms[AgeGroup(i)] = 1 + model.parameters.TransmissionProbabilityOnContact[AgeGroup(i)] = rho + model.parameters.RecoveredPerInfectedNoSymptoms[AgeGroup(i)] = muCR + model.parameters.RiskOfInfectionFromSymptomatic[AgeGroup(i)] = 0.25 + model.parameters.SeverePerInfectedSymptoms[AgeGroup(i)] = muHI + model.parameters.CriticalPerSevere[AgeGroup(i)] = muUH + model.parameters.DeathsPerCritical[AgeGroup(i)] = muDU + # twice the value of RiskOfInfectionFromSymptomatic + model.parameters.MaxRiskOfInfectionFromSymptomatic[AgeGroup(i)] = 0.5 + # StartDay is the n-th day of the year + model.parameters.StartDay = start_date.day_in_year + + +def set_contact_matrices(model, data_dir, num_groups=6): + """Setting the contact matrices for a model + + :param model: memilio ODE-model, whose contact matrices should be modified. + :param data_dir: directory, where the contact data is stored (should contain folder "contacts") + :param num_groups: Number of age groups considered + + """ + + contact_matrices = mio.ContactMatrixGroup( + len(list(Location)), num_groups) + locations = ["home", "school_pf_eig", "work", "other"] + + # Loading contact matrices for each location from .txt file + for i, location in enumerate(locations): + baseline_file = os.path.join( + data_dir, "Germany", "contacts", "baseline_" + location + ".txt") + + contact_matrices[i] = mio.ContactMatrix( + mio.read_mobility_plain(baseline_file), + ) + model.parameters.ContactPatterns.cont_freq_mat = contact_matrices + + +def extrapolate_data(start_date, num_days, data_dir): + model = Model(6) + set_covid_parameters(model, start_date) + set_contact_matrices(model, data_dir) + + graph = ModelGraph() + + scaling_factor_infected = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + scaling_factor_icu = 1.0 + tnt_capacity_factor = 7.5 / 100000. + + # Path containing the population data + data_dir_Germany = os.path.join(data_dir, "Germany") + pydata_dir = os.path.join(data_dir_Germany, "pydata") + + path_population_data = os.path.join(pydata_dir, + "county_current_population.json") + + set_nodes( + model.parameters, + start_date, start_date + num_days, + pydata_dir, + path_population_data, True, graph, scaling_factor_infected, + scaling_factor_icu, tnt_capacity_factor, num_days, True) + + +def read_results_h5(path, group_key='Total'): + with h5py.File(path, 'r') as f: + keys = list(f.keys()) + res = {} + for i, key in enumerate(keys): + group = f[key] + total = group[group_key][()] + # remove confirmed compartments + sum_inf_no_symp = np.sum(total[:, [2, 3]], axis=1) + sum_inf_symp = np.sum(total[:, [4, 5]], axis=1) + total[:, 2] = sum_inf_no_symp + total[:, 4] = sum_inf_symp + total = np.delete(total, [3, 5], axis=1) + res[key] = total + return res + + +def get_ground_truth_data(start_date, num_days, data_dir, create_new=False): + path_rki_h5 = os.path.join(data_dir, "Germany", "pydata", "Results_rki.h5") + if not os.path.isfile(path_rki_h5) or create_new: + print("Generating real data from C++ backend...") + extrapolate_data(start_date, num_days, data_dir) + print("Data generation complete.") + + # The GNN expects age-stratified input. + num_age_groups = 6 + all_age_data_list = [] + group_keys = ['Group1', 'Group2', 'Group3', 'Group4', 'Group5', 'Group6'] + for age in range(num_age_groups): + age_data_dict = read_results_h5(path_rki_h5, group_keys[age]) + age_data_np = np.array(list(age_data_dict.values())) + all_age_data_list.append(age_data_np) + + # Combine age groups to get shape (num_nodes, timesteps, num_features=48) + all_age_data = np.stack(all_age_data_list, axis=-1) + num_nodes, timesteps, _, _ = all_age_data.shape + ground_truth_all_nodes_np = np.reshape( + all_age_data.transpose(0, 1, 3, 2), (num_nodes, timesteps, -1)) + + return ground_truth_all_nodes_np + + +def extrapolate_ground_truth_data(data_dir, num_days=360, num_input_days=0): + start_dates = [ + mio.Date(2020, 6, 1) + # mio.Date(2020, 7, 1), + # mio.Date(2020, 8, 1), + # mio.Date(2020, 9, 1), + # mio.Date(2020, 10, 1), + # mio.Date(2020, 11, 1) + # mio.Date(2020, 12, 1), + ] + + for start_date in start_dates: + ground_truth_np = get_ground_truth_data( + start_date, num_days, data_dir, create_new=True) + path_rki_h5 = "/localdata1/hege_mn/memilio/data/Germany/pydata/Results_rki.h5" + num_age_groups = 6 + all_age_data_list = [] + group_keys = ['Group1', 'Group2', 'Group3', 'Group4', 'Group5', 'Group6'] + for age in range(num_age_groups): + age_data_dict = read_results_h5(path_rki_h5, group_keys[age]) + age_data_np = np.array(list(age_data_dict.values())) + all_age_data_list.append(age_data_np) + # Combine age groups to get shape (num_nodes, timesteps, num_features=48) + all_age_data = np.stack(all_age_data_list, axis=-1) + + # Save the ground truth data + output_path = "/localdata1/hege_mn/memilio/data/Germany/pydata/ground_truth_all_nodes.pickle" + with open(output_path, 'wb') as f: + pd.to_pickle(all_age_data, f) + print(f"Ground truth data saved to {output_path}") + + +def main(): + cwd = os.getcwd() + start_date = mio.Date(2020, 8, 1) + num_days = 360 + num_input_days = 5 + number_of_channels = 512 + data_dir = os.path.join(cwd, "data") + + # path_weights = '/localdata1/hege_mn/memilio/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/saved_weights_paper/' + # path_model_weights = os.path.join( + # path_weights, "GNN_30days_nodeswithvariance_1k_3Damp_test2.pickle") + + # compare_and_predict( + # start_date, + # num_days, + # data_dir, + # path_model_weights, + # number_of_channels, + # num_input_days + # ) + + extrapolate_ground_truth_data( + data_dir, + num_days, + num_input_days=0 + ) + + +if __name__ == "__main__": + main() From 0ceefd4395cc92166a8b2bf2084915ab9d5ed47e Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 11 Aug 2025 13:11:11 +0200 Subject: [PATCH 02/41] reducing extrapolate gnn to its core functionality, adding comments in data_generation --- .../surrogatemodel/GNN/data_generation.py | 21 +- .../GNN/extrapolate_data_gnn.py | 182 +++++++----------- 2 files changed, 79 insertions(+), 124 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index ee6a97edb3..4367d054fc 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -48,6 +48,7 @@ class Location(Enum): Other = 3 +# Define the start and the latest end date for the simulation start_date = mio.Date(2019, 1, 1) end_date = mio.Date(2021, 12, 31) @@ -84,7 +85,7 @@ def set_covid_parameters(model, num_groups=6): model.parameters.TimeInfectedSevere[AgeGroup(i)] = th model.parameters.TimeInfectedCritical[AgeGroup(i)] = tu - # Compartment transition propabilities + # Compartment transition probabilities model.parameters.RelativeTransmissionNoSymptoms[AgeGroup(i)] = 1 model.parameters.TransmissionProbabilityOnContact[AgeGroup(i)] = rho model.parameters.RecoveredPerInfectedNoSymptoms[AgeGroup(i)] = muCR @@ -190,23 +191,25 @@ def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_ start_date_num = start_date.day_in_year - min_date_num # Load the ground truth data - with open("/localdata1/hege_mn/memilio/data/Germany/pydata/ground_truth_all_nodes.pickle", 'rb') as f: - ground_truth_all_nodes_np = pickle.load(f) + pydata_dir = os.path.join(data_dir, "Germany", "pydata") + ground_truth_dir = os.path.join( + pydata_dir, "ground_truth_all_nodes.pickle") + with open(ground_truth_dir, 'rb') as f: + ground_truth_all_nodes = pickle.load(f) # Initialize model for each node, using the population data and sampling the number of # individuals in the different compartments for node_indx in range(graph.num_nodes): model = graph.get_node(node_indx).property - data = ground_truth_all_nodes_np[node_indx][start_date_num] + data = ground_truth_all_nodes[node_indx][start_date_num] - # Set parameters - # TODO: Put This in the draw_sample function in the ParameterStudy # Iterate over the different age groups for i in range(num_groups): age_group = AgeGroup(i) pop_age_group = model.populations.get_group_total_AgeGroup( age_group) + # Generating valid, noisy configuration of the compartments valid_configuration = False while valid_configuration is False: comp_data = data[:, i] @@ -245,7 +248,6 @@ def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_ ) * np.float16(factor) model.parameters.ContactPatterns.cont_freq_mat.add_damping(Damping( coeffs=(damping), t=day, level=0, type=0)) -# damping matrices and damping_coefficients werden überschrieben -> am Ende nur von letztem Knoten übergeben ??? damped_matrices.append(model.parameters.ContactPatterns.cont_freq_mat.get_matrix_at( day+1)) damping_coefficients.append(damping) @@ -303,8 +305,9 @@ def generate_data( # Setting basic parameter num_groups = 6 - mobility_dir = "/localdata1/hege_mn/memilio/data/Germany/mobility/commuter_mobility_2022.txt" + mobility_dir = data_dir + "/Germany/mobility/commuter_mobility_2022.txt" graph = get_graph(num_groups, data_dir, mobility_dir) + # Define possible start dates for the simulation start_dates = [ mio.Date(2020, 6, 1), mio.Date(2020, 7, 1), @@ -316,7 +319,7 @@ def generate_data( ] # show progess in terminal for longer runs - # Due to the random structure, theres currently no need to shuffle the data + # Due to the random structure, there is currently no need to shuffle the data bar = Bar('Number of Runs done', max=num_runs) times = [] diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py index 3a6f84b9b0..4a9ac15143 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py @@ -1,103 +1,56 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Henrik Zunker +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + import os import numpy as np -import pickle import h5py import pandas as pd -import matplotlib.pyplot as plt import memilio.simulation as mio -import tensorflow as tf +from memilio.simulation import (AgeGroup) -from enum import Enum -from memilio.simulation import (AgeGroup, LogLevel, set_log_level, Damping) from memilio.simulation.osecir import ( InfectionState, Model, ModelGraph, set_nodes) from memilio.surrogatemodel.GNN.GNN_utils import remove_confirmed_compartments - - -class Location(Enum): - Home = 0 - School = 1 - Work = 2 - Other = 3 - - -def set_covid_parameters(model, start_date, num_groups=6): - """Setting COVID-parameters for the different age groups. - - :param model: memilio model, whose parameters should be clarified - :param num_groups: Number of age groups - """ - - # age specific parameters - TransmissionProbabilityOnContact = [0.03, 0.06, 0.06, 0.06, 0.09, 0.175] - RecoveredPerInfectedNoSymptoms = [0.25, 0.25, 0.2, 0.2, 0.2, 0.2] - SeverePerInfectedSymptoms = [0.0075, 0.0075, 0.019, 0.0615, 0.165, 0.225] - CriticalPerSevere = [0.075, 0.075, 0.075, 0.15, 0.3, 0.4] - DeathsPerCritical = [0.05, 0.05, 0.14, 0.14, 0.4, 0.6] - - TimeInfectedNoSymptoms = [2.74, 2.74, 2.565, 2.565, 2.565, 2.565] - TimeInfectedSymptoms = [7.02625, 7.02625, - 7.0665, 6.9385, 6.835, 6.775] - TimeInfectedSevere = [5, 5, 5.925, 7.55, 8.5, 11] - TimeInfectedCritical = [6.95, 6.95, 6.86, 17.36, 17.1, 11.6] - - for i, rho, muCR, muHI, muUH, muDU, tc, ti, th, tu in zip(range(num_groups), - TransmissionProbabilityOnContact, RecoveredPerInfectedNoSymptoms, - SeverePerInfectedSymptoms, CriticalPerSevere, DeathsPerCritical, - TimeInfectedNoSymptoms, TimeInfectedSymptoms, - TimeInfectedSevere, TimeInfectedCritical): - # Compartment transition duration - model.parameters.TimeExposed[AgeGroup(i)] = 3.335 - model.parameters.TimeInfectedNoSymptoms[AgeGroup(i)] = tc - model.parameters.TimeInfectedSymptoms[AgeGroup(i)] = ti - model.parameters.TimeInfectedSevere[AgeGroup(i)] = th - model.parameters.TimeInfectedCritical[AgeGroup(i)] = tu - - # Compartment transition propabilities - model.parameters.RelativeTransmissionNoSymptoms[AgeGroup(i)] = 1 - model.parameters.TransmissionProbabilityOnContact[AgeGroup(i)] = rho - model.parameters.RecoveredPerInfectedNoSymptoms[AgeGroup(i)] = muCR - model.parameters.RiskOfInfectionFromSymptomatic[AgeGroup(i)] = 0.25 - model.parameters.SeverePerInfectedSymptoms[AgeGroup(i)] = muHI - model.parameters.CriticalPerSevere[AgeGroup(i)] = muUH - model.parameters.DeathsPerCritical[AgeGroup(i)] = muDU - # twice the value of RiskOfInfectionFromSymptomatic - model.parameters.MaxRiskOfInfectionFromSymptomatic[AgeGroup(i)] = 0.5 - # StartDay is the n-th day of the year - model.parameters.StartDay = start_date.day_in_year - - -def set_contact_matrices(model, data_dir, num_groups=6): - """Setting the contact matrices for a model - - :param model: memilio ODE-model, whose contact matrices should be modified. - :param data_dir: directory, where the contact data is stored (should contain folder "contacts") - :param num_groups: Number of age groups considered - - """ - - contact_matrices = mio.ContactMatrixGroup( - len(list(Location)), num_groups) - locations = ["home", "school_pf_eig", "work", "other"] - - # Loading contact matrices for each location from .txt file - for i, location in enumerate(locations): - baseline_file = os.path.join( - data_dir, "Germany", "contacts", "baseline_" + location + ".txt") - - contact_matrices[i] = mio.ContactMatrix( - mio.read_mobility_plain(baseline_file), - ) - model.parameters.ContactPatterns.cont_freq_mat = contact_matrices +from memilio.surrogatemodel.GNN.data_generation import ( + set_contact_matrices, set_covid_parameters, Location) def extrapolate_data(start_date, num_days, data_dir): + ''' + Extrapolate data using the C++ backend. + This function sets up an empty model and calls the C++ backend to proceed the real world + data. + :param start_date: The date to start the simulation from. + :param num_days: The number of days to simulate. + :param data_dir: The directory where the data is stored. + :return: None + ''' + # Set up the model model = Model(6) - set_covid_parameters(model, start_date) + model.parameters.StartDay = start_date.day_in_year + set_covid_parameters(model) set_contact_matrices(model, data_dir) graph = ModelGraph() + # Set the parameters for the nodes scaling_factor_infected = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] scaling_factor_icu = 1.0 tnt_capacity_factor = 7.5 / 100000. @@ -108,7 +61,7 @@ def extrapolate_data(start_date, num_days, data_dir): path_population_data = os.path.join(pydata_dir, "county_current_population.json") - + # Set the nodes in the model set_nodes( model.parameters, start_date, start_date + num_days, @@ -118,6 +71,11 @@ def extrapolate_data(start_date, num_days, data_dir): def read_results_h5(path, group_key='Total'): + ''' Reads results from an HDF5 file and processes the data. + :param path: Path to the HDF5 file. + :param group_key: The key for the group to read from the HDF5 file. + :return: A dictionary with processed results. + ''' with h5py.File(path, 'r') as f: keys = list(f.keys()) res = {} @@ -135,13 +93,22 @@ def read_results_h5(path, group_key='Total'): def get_ground_truth_data(start_date, num_days, data_dir, create_new=False): + ''' + Generates ground truth data for the specified start date and number of days. + :param start_date: The date to start the simulation from. + :param num_days: The number of days to simulate. + :param data_dir: The directory where the data is stored. + :param create_new: If True, generates new data; otherwise, reads existing data. + :return: A numpy array containing the ground truth data. + ''' + # Check if the path exists and create new data if needed path_rki_h5 = os.path.join(data_dir, "Germany", "pydata", "Results_rki.h5") if not os.path.isfile(path_rki_h5) or create_new: print("Generating real data from C++ backend...") extrapolate_data(start_date, num_days, data_dir) print("Data generation complete.") - # The GNN expects age-stratified input. + # age-stratified input. num_age_groups = 6 all_age_data_list = [] group_keys = ['Group1', 'Group2', 'Group3', 'Group4', 'Group5', 'Group6'] @@ -150,7 +117,6 @@ def get_ground_truth_data(start_date, num_days, data_dir, create_new=False): age_data_np = np.array(list(age_data_dict.values())) all_age_data_list.append(age_data_np) - # Combine age groups to get shape (num_nodes, timesteps, num_features=48) all_age_data = np.stack(all_age_data_list, axis=-1) num_nodes, timesteps, _, _ = all_age_data.shape ground_truth_all_nodes_np = np.reshape( @@ -159,21 +125,25 @@ def get_ground_truth_data(start_date, num_days, data_dir, create_new=False): return ground_truth_all_nodes_np -def extrapolate_ground_truth_data(data_dir, num_days=360, num_input_days=0): +def extrapolate_ground_truth_data(data_dir, num_days=360): + ''' + Extrapolates ground truth data for the specified number of days. + :param data_dir: The directory where the data is stored. + :param num_days: The number of days to simulate. + :return: None + ''' + cwd = os.getcwd() + data_dir = os.path.join(cwd, data_dir) start_dates = [ mio.Date(2020, 6, 1) - # mio.Date(2020, 7, 1), - # mio.Date(2020, 8, 1), - # mio.Date(2020, 9, 1), - # mio.Date(2020, 10, 1), - # mio.Date(2020, 11, 1) - # mio.Date(2020, 12, 1), ] for start_date in start_dates: - ground_truth_np = get_ground_truth_data( + get_ground_truth_data( start_date, num_days, data_dir, create_new=True) - path_rki_h5 = "/localdata1/hege_mn/memilio/data/Germany/pydata/Results_rki.h5" + print(f"Ground truth data for {start_date} generated.") + + path_rki_h5 = data_dir + "/Germany/pydata/Results_rki.h5" num_age_groups = 6 all_age_data_list = [] group_keys = ['Group1', 'Group2', 'Group3', 'Group4', 'Group5', 'Group6'] @@ -181,11 +151,10 @@ def extrapolate_ground_truth_data(data_dir, num_days=360, num_input_days=0): age_data_dict = read_results_h5(path_rki_h5, group_keys[age]) age_data_np = np.array(list(age_data_dict.values())) all_age_data_list.append(age_data_np) - # Combine age groups to get shape (num_nodes, timesteps, num_features=48) all_age_data = np.stack(all_age_data_list, axis=-1) # Save the ground truth data - output_path = "/localdata1/hege_mn/memilio/data/Germany/pydata/ground_truth_all_nodes.pickle" + output_path = data_dir + "/Germany/pydata/ground_truth_all_nodes.pickle" with open(output_path, 'wb') as f: pd.to_pickle(all_age_data, f) print(f"Ground truth data saved to {output_path}") @@ -193,29 +162,12 @@ def extrapolate_ground_truth_data(data_dir, num_days=360, num_input_days=0): def main(): cwd = os.getcwd() - start_date = mio.Date(2020, 8, 1) num_days = 360 - num_input_days = 5 - number_of_channels = 512 data_dir = os.path.join(cwd, "data") - # path_weights = '/localdata1/hege_mn/memilio/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/saved_weights_paper/' - # path_model_weights = os.path.join( - # path_weights, "GNN_30days_nodeswithvariance_1k_3Damp_test2.pickle") - - # compare_and_predict( - # start_date, - # num_days, - # data_dir, - # path_model_weights, - # number_of_channels, - # num_input_days - # ) - extrapolate_ground_truth_data( data_dir, - num_days, - num_input_days=0 + num_days ) From fde4b3b3cba0eab635006c035cb7276918f59dbf Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 11 Aug 2025 13:17:36 +0200 Subject: [PATCH 03/41] Add description of workflow --- .../memilio/surrogatemodel/GNN/README.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md new file mode 100644 index 0000000000..6c1dd3e54d --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md @@ -0,0 +1,5 @@ +GNN package +=============== +1) Proceed the real world data by running extrapolate_data_gnn.py. The data is saved as "ground_truth_all_nodes.pickle" in +the pydata folder +2) Generate data using data_generation.py. The model is initialized using the real world data + some noise. \ No newline at end of file From ba90b765c84e08778923d8c97e1af3d9064cfef5 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Mon, 11 Aug 2025 13:52:33 +0200 Subject: [PATCH 04/41] init file gnn --- .../memilio/surrogatemodel/GNN/__init__.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/__init__.py diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/__init__.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/__init__.py new file mode 100644 index 0000000000..d97497a6ee --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/__init__.py @@ -0,0 +1,23 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Agatha Schmidt, Henrik Zunker +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + +""" +A surrogate model for a spatial resolved SECIR model. +""" From 3c9bf9bf5643e066e93e890df18bcdcc8c25e5e9 Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 18 Aug 2025 14:10:24 +0200 Subject: [PATCH 05/41] [ci skip] Updated data generation and extrapolation, code for training and evaluation, generating network architectures (draft) --- .../surrogatemodel/GNN/data_generation.py | 178 ++++++++-- .../surrogatemodel/GNN/evaluate_and_train.py | 303 ++++++++++++++++++ .../GNN/extrapolate_data_gnn.py | 22 +- .../GNN/network_architectures.py | 87 +++++ 4 files changed, 561 insertions(+), 29 deletions(-) create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index 4367d054fc..9261356720 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -167,7 +167,7 @@ def get_graph(num_groups, data_dir, mobility_directory): return graph -def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_groups=6, start_date=mio.Date(2020, 6, 1)): +def run_secir_groups_simulation1(days, damping_days, damping_factors, graph, num_groups=6, start_date=mio.Date(2020, 6, 1)): """ Uses an ODE SECIR model allowing for asymptomatic infection with 6 different age groups. The model is not stratified by region. Virus-specific parameters are fixed and initial number of person @@ -214,25 +214,142 @@ def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_ while valid_configuration is False: comp_data = data[:, i] ratios = np.random.uniform(0.8, 1.2, size=8) - modified_data = comp_data * ratios - if np.sum(modified_data[1:]) < pop_age_group: + init_data = comp_data * ratios + if np.sum(init_data[1:]) < pop_age_group: + valid_configuration = True + + # Set the populations for the different compartments + model.populations[age_group, Index_InfectionState( + InfectionState.Exposed)] = init_data[1] + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedNoSymptoms)] = init_data[2] + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedSymptoms)] = init_data[3] + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedSevere)] = init_data[4] + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedCritical)] = init_data[5] + model.populations[age_group, Index_InfectionState( + InfectionState.Recovered)] = init_data[6] + model.populations[age_group, Index_InfectionState( + InfectionState.Dead)] = init_data[7] + model.populations.set_difference_from_group_total_AgeGroup(( + age_group, InfectionState.Susceptible), pop_age_group) + + # Introduce the damping information, in general dampings can be local, but till now the code just allows global dampings + damped_matrices = [] + damping_coefficients = [] + + for i in np.arange(len(damping_days)): + day = damping_days[i] + factor = damping_factors[i] + + damping = np.ones((num_groups, num_groups) + ) * np.float16(factor) + model.parameters.ContactPatterns.cont_freq_mat.add_damping(Damping( + coeffs=(damping), t=day, level=0, type=0)) + damped_matrices.append(model.parameters.ContactPatterns.cont_freq_mat.get_matrix_at( + day+1)) + damping_coefficients.append(damping) + + # Apply mathematical constraints to parameters + model.apply_constraints() + + # set model to graph + graph.get_node(node_indx).property.populations = model.populations + + # Start simulation + study = ParameterStudy(graph, 0, days, dt=0.5, num_runs=1) + start_time = time.perf_counter() + study.run() + runtime = time.perf_counter() - start_time + + graph_run = study.run()[0] + results = interpolate_simulation_result(graph_run) + + for result_indx in range(len(results)): + results[result_indx] = remove_confirmed_compartments( + np.asarray(results[result_indx]), num_groups) + + dataset_entry = copy.deepcopy(results) + + return dataset_entry, damped_matrices, damping_coefficients, runtime + + +def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_groups=6, start_date=mio.Date(2020, 6, 1)): + """ Uses an ODE SECIR model allowing for asymptomatic infection with 6 + different age groups. The model is not stratified by region. + Virus-specific parameters are fixed and initial number of person + in the particular infection states are chosen randomly from defined ranges. + + :param days: Number of days simulated within a single run. + :param damping_days: Days, where a damping is applied + :param damping_factors: damping factors associated to the damping days + :param graph: Graph initialized for the start_date with the population data which + is sampled during the run. + :param num_groups: Number of age groups considered in the simulation + :param start_date: Date, when the simulation starts + :returns: List containing the populations in each compartment used to initialize + the run. + """ + if len(damping_days) != len(damping_factors): + raise ValueError("Length of damping_days and damping_factors differ!") + + min_date = mio.Date(2020, 6, 1) + min_date_num = min_date.day_in_year + start_date_num = start_date.day_in_year - min_date_num + + # Load the ground truth data + pydata_dir = os.path.join(data_dir, "Germany", "pydata") + upper_bound_dir = os.path.join( + pydata_dir, "ground_truth_upper_bound.pickle") + lower_bound_dir = os.path.join( + pydata_dir, "ground_truth_lower_bound.pickle") + with open(upper_bound_dir, 'rb') as f: + ground_truth_upper_bound = pickle.load(f) + + with open(lower_bound_dir, 'rb') as f: + ground_truth_lower_bound = pickle.load(f) + + # Initialize model for each node, using the population data and sampling the number of + # individuals in the different compartments + for node_indx in range(graph.num_nodes): + model = graph.get_node(node_indx).property + max_data = ground_truth_upper_bound[node_indx] + min_data = ground_truth_lower_bound[node_indx] + + # Iterate over the different age groups + for i in range(num_groups): + age_group = AgeGroup(i) + pop_age_group = model.populations.get_group_total_AgeGroup( + age_group) + + # Generating valid, noisy configuration of the compartments + valid_configuration = False + while valid_configuration is False: + init_data = np.asarray([0 for _ in range(8)]) + max_val = max_data[:, i] + min_val = min_data[:, i] + for j in range(1, 8): + init_data[j] = random.uniform(min_val[j], max_val[j]) + if np.sum(init_data[1:]) < pop_age_group: valid_configuration = True # Set the populations for the different compartments model.populations[age_group, Index_InfectionState( - InfectionState.Exposed)] = modified_data[1] + InfectionState.Exposed)] = init_data[1] model.populations[age_group, Index_InfectionState( - InfectionState.InfectedNoSymptoms)] = modified_data[2] + InfectionState.InfectedNoSymptoms)] = init_data[2] model.populations[age_group, Index_InfectionState( - InfectionState.InfectedSymptoms)] = modified_data[3] + InfectionState.InfectedSymptoms)] = init_data[3] model.populations[age_group, Index_InfectionState( - InfectionState.InfectedSevere)] = modified_data[4] + InfectionState.InfectedSevere)] = init_data[4] model.populations[age_group, Index_InfectionState( - InfectionState.InfectedCritical)] = modified_data[5] + InfectionState.InfectedCritical)] = init_data[5] model.populations[age_group, Index_InfectionState( - InfectionState.Recovered)] = modified_data[6] + InfectionState.Recovered)] = init_data[6] model.populations[age_group, Index_InfectionState( - InfectionState.Dead)] = modified_data[7] + InfectionState.Dead)] = init_data[7] model.populations.set_difference_from_group_total_AgeGroup(( age_group, InfectionState.Susceptible), pop_age_group) @@ -308,15 +425,15 @@ def generate_data( mobility_dir = data_dir + "/Germany/mobility/commuter_mobility_2022.txt" graph = get_graph(num_groups, data_dir, mobility_dir) # Define possible start dates for the simulation - start_dates = [ - mio.Date(2020, 6, 1), - mio.Date(2020, 7, 1), - mio.Date(2020, 8, 1), - mio.Date(2020, 9, 1), - mio.Date(2020, 10, 1), - mio.Date(2020, 11, 1), - mio.Date(2020, 12, 1) - ] + # start_dates = [ + # mio.Date(2020, 6, 1), + # mio.Date(2020, 7, 1), + # mio.Date(2020, 8, 1), + # mio.Date(2020, 9, 1), + # mio.Date(2020, 10, 1), + # mio.Date(2020, 11, 1), + # mio.Date(2020, 12, 1) + # ] # show progess in terminal for longer runs # Due to the random structure, there is currently no need to shuffle the data @@ -324,11 +441,15 @@ def generate_data( times = [] for i in range(0, num_runs): - start_date = start_dates[i % len(start_dates)] + # start_date = start_dates[i % len(start_dates)] # Generate random damping days and damping factors - damping_days, damping_factors = dampings.generate_dampings( - days, max_number_damping, method=damping_method, min_distance=2, - min_damping_day=2) + if max_number_damping > 0: + damping_days, damping_factors = dampings.generate_dampings( + days, max_number_damping, method=damping_method, min_distance=2, + min_damping_day=2) + else: + damping_days = [] + damping_factors = [] # Run simulation data_run, damped_matrices, damping_coefficients, t_run = run_secir_groups_simulation( days, damping_days, damping_factors, graph, num_groups, start_date) @@ -355,8 +476,8 @@ def generate_data( if transform: inputs, labels = scale_data(data) - all_data = {"inputs": scaled_inputs, - "labels": scaled_labels, + all_data = {"inputs": inputs, + "labels": labels, "damping_day": data["damping_days"], "contact_matrix": data["contact_matrix"], "damping_coeff": data["damping_factors"] @@ -388,8 +509,8 @@ def generate_data( # data_dir = os.path.join(os.getcwd(), 'data') data_dir = "/localdata1/hege_mn/memilio/data" input_width = 5 - number_of_dampings = 3 - num_runs = 30 + number_of_dampings = 0 + num_runs = 100 label_width_list = [30] random.seed(10) @@ -399,5 +520,6 @@ def generate_data( path=path_output, input_width=input_width, label_width=label_width, - save_data=False, + save_data=True, + damping_method="active", max_number_damping=number_of_dampings) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py new file mode 100644 index 0000000000..4759ad194f --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -0,0 +1,303 @@ +import os +import pickle +import spektral +import time +import pandas as pd +import numpy as np +from sklearn.preprocessing import FunctionTransformer + + +from tensorflow.keras.optimizers import Adam +from tensorflow.keras.losses import MeanAbsolutePercentageError +from tensorflow.keras.models import Model +import tensorflow.keras.initializers as initializers + +import copy +import tensorflow as tf +import memilio.surrogatemodel.GNN.network_architectures as network_architectures +from memilio.surrogatemodel.utils.helper_functions import (calc_split_index) + +from spektral.transforms.normalize_adj import NormalizeAdj +from spektral.data import Dataset, DisjointLoader, Graph, Loader, BatchLoader, MixedLoader +from spektral.layers import GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv +from spektral.utils.convolution import normalized_laplacian, rescale_laplacian + + +def create_dataset(path_cases, path_mobility): + """ + Create a dataset from the given file path. + """ + file = open( + path_cases, 'rb') + data = pickle.load(file) + + number_of_nodes = 400 + + inputs = data['inputs'] + labels = data['labels'] + + len_dataset = len(inputs) + + shape_input_flat = np.asarray( + inputs).shape[1]*np.asarray(inputs).shape[2] + shape_labels_flat = np.asarray( + labels).shape[1]*np.asarray(labels).shape[2] + + new_inputs = np.asarray( + inputs.transpose(0, 3, 1, 2)).reshape( + len_dataset, number_of_nodes, shape_input_flat) + new_labels = np.asarray(labels.transpose(0, 3, 1, 2)).reshape( + len_dataset, number_of_nodes, shape_labels_flat) + + commuter_file = open(os.path.join( + path_mobility, 'commuter_mobility_2022.txt'), 'rb') + commuter_data = pd.read_csv(commuter_file, sep=" ", header=None) + sub_matrix = commuter_data.iloc[:number_of_nodes, 0:number_of_nodes] + + adjacency_matrix = np.asarray(sub_matrix) + + adjacency_matrix[adjacency_matrix > 0] = 1 # make the adjacency binary + node_features = new_inputs + + node_labels = new_labels + + class MyDataset(spektral.data.dataset.Dataset): + def read(self): + self.a = adjacency_matrix + + # self.a = normalized_adjacency(adjacency_matrix) + # self.a = rescale_laplacian(normalized_laplacian(adjacency_matrix)) + + return [spektral.data.Graph(x=x, y=y) for x, y in zip(node_features, node_labels)] + + super().__init__(**kwargs) + + # data = MyDataset() + data = MyDataset() + + return data + + +def train_step(inputs, target, loss_fn, model, optimizer): + with tf.GradientTape() as tape: + predictions = model(inputs, training=True) + loss = loss_fn(target, predictions) + \ + sum(model.losses) # Add regularization losses + gradients = tape.gradient(loss, model.trainable_variables) + optimizer.apply_gradients(zip(gradients, model.trainable_variables)) + acc = tf.reduce_mean(loss_fn(target, predictions)) + return loss, acc + + +def evaluate(loader, model, loss_fn, retransform=False): + output = [] + step = 0 + while step < loader.steps_per_epoch: + step += 1 + inputs, target = loader.__next__() + pred = model(inputs, training=False) + if retransform: + target = np.expm1(target) + pred = np.expm1(pred) + # Calculate loss and metrics + outs = ( + loss_fn(target, pred), + tf.reduce_mean(loss_fn(target, pred)), + len(target), # Keep track of batch size + ) + output.append(outs) + if step == loader.steps_per_epoch: + output = np.array(output) + return np.average(output[:, :-1], 0, weights=output[:, -1]) + + +def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_patience, save_name): + n = len(data) + n_train, n_valid, n_test = calc_split_index( + n, split_train=0.7, split_valid=0.2, split_test=0.1) + + # Split data into train, validation, and test sets + train_data, valid_data, test_data = data[:n_train], data[n_train:n_train + + n_valid], data[n_train+n_valid:] + + # Define Data Loaders + loader_tr = MixedLoader( + train_data, batch_size=batch_size, epochs=epochs, shuffle=False) + loader_val = MixedLoader( + valid_data, batch_size=batch_size, shuffle=False) + loader_test = MixedLoader( + test_data, batch_size=n_test, shuffle=False) + + df = pd.DataFrame(columns=[ + "train_loss", "val_loss", "test_loss", + "test_loss_orig", "training_time", + "loss_history", "val_loss_history"]) + + test_scores = [] + test_scores_r = [] + train_losses = [] + val_losses = [] + + losses_history_all = [] + val_losses_history_all = [] + + best_val_loss = np.inf + best_weights = None + patience = es_patience + results = [] + losses_history = [] + val_losses_history = [] + + start = time.perf_counter() + step = 0 + epoch = 0 + for batch in loader_tr: + step += 1 + loss, acc = train_step(*batch, loss_fn, model, optimizer) + results.append((loss, acc)) + + # Compute validation loss and accuracy + if step == loader_tr.steps_per_epoch: + step = 0 + epoch += 1 + val_loss, val_acc = evaluate(loader_val, model, loss_fn) + print( + "Ep. {} - Loss: {:.3f} - Acc: {:.3f} - Val loss: {:.3f} - Val acc: {:.3f}".format( + epoch, *np.mean(results, 0), val_loss, val_acc + ) + ) + # Check if loss improved for early stopping + if val_loss < best_val_loss: + best_val_loss = val_loss + patience = es_patience + print(f"New best val_loss {val_loss:.3f}") + best_weights = model.get_weights() + else: + patience -= 1 + if patience == 0: + print( + "Early stopping (best val_loss: {})".format( + best_val_loss)) + break + results = [] + losses_history.append(loss) + val_losses_history.append(val_loss) + ################################################################################ + # Evaluate model + ################################################################################ + model.set_weights(best_weights) # Load best model + test_loss, test_acc = evaluate(loader_test, model, loss_fn) + test_loss_r, test_acc_r = evaluate( + loader_test, model, loss_fn, retransform=True) + + print( + "Done. Test loss: {:.4f}. Test acc: {:.2f}".format( + test_loss, test_acc)) + test_scores.append(test_loss) + test_scores_r.append(test_loss_r) + train_losses.append(np.asarray(losses_history).min()) + val_losses.append(np.asarray(val_losses_history).min()) + losses_history_all.append(np.asarray(losses_history)) + val_losses_history_all.append(val_losses_history) + + elapsed = time.perf_counter() - start + + # print out stats + print(f"Best train losses: {train_losses} ") + print(f"Best validation losses: {val_losses}") + print(f"Test values: {test_scores}") + print("--------------------------------------------") + print(f"Train Score:{np.mean(train_losses)}") + print(f"Validation Score:{np.mean(val_losses)}") + print(f"Test Score (log): {np.mean(test_scores)}") + print(f"Test Score (orig.): {np.mean(test_scores_r)}") + + print(f"Time for training: {elapsed:.4f} seconds") + print(f"Time for training: {elapsed/60:.4f} minutes") + + +# save df + df.loc[len(df.index)] = [ # layer_name, number_of_layer, channels, + np.mean(train_losses), + np.mean(val_losses), + np.mean(test_scores), + np.mean(test_scores_r), + (elapsed / 60), + [losses_history_all], + [val_losses_history_all]] + print(df) + + # Ensure that save_name has the .pickle extension + if not save_name.endswith('.pickle'): + save_name += '.pickle' + + # Save best weights as pickle + path = os.path.dirname(os.path.realpath(__file__)) + file_path_w = os.path.join(path, 'saved_weights') + # Ensure the directory exists + if not os.path.isdir(file_path_w): + os.mkdir(file_path_w) + + # Construct the full file path by joining the directory with save_name + file_path_w = os.path.join(file_path_w, save_name) + + # Save the weights to the file + with open(file_path_w, 'wb') as f: + pickle.dump(best_weights, f) + + file_path_df = os.path.join( + os.path.dirname( + os.path.realpath(os.path.dirname(os.path.realpath(path)))), + 'model_evaluations_paper') + if not os.path.isdir(file_path_df): + os.mkdir(file_path_df) + file_path_df = file_path_df+save_name.replace('.pickle', '.csv') + df.to_csv(file_path_df) + + +if __name__ == "__main__": + start_hyper = time.perf_counter() + epochs = 10 + batch_size = 2 + es_patience = 10 + optimizer = Adam(learning_rate=0.001) + loss_fn = MeanAbsolutePercentageError() + + # Generate the Dataset + path_cases = "/localdata1/hege_mn/memilio/saves/GNN_data_30days_3dampings_classic5.pickle" + path_mobility = '/localdata1/hege_mn/memilio/data/Germany/mobility' + data = create_dataset(path_cases, path_mobility) + + # Define the model architecture + def transform_a(adjacency_matrix): + a = adjacency_matrix.numpy() + a = rescale_laplacian(normalized_laplacian(a)) + return tf.convert_to_tensor(a, dtype=tf.float32) + + layer_types = [ + # Dense layer (only uses x) + lambda: ARMAConv(512, activation='elu', + kernel_initializer=initializers.GlorotUniform(seed=None)) + ] + num_repeat = [7] + + model_class = network_architectures.generate_model_class( + "ARMA", layer_types, num_repeat, num_output=1440, transform=transform_a) + + model = model_class() + + # early stopping patience + # name for df + save_name = 'GNN_30days' # name for model + # for param in parameters: + + train_and_evaluate(data, batch_size, epochs, model, + loss_fn, optimizer, es_patience, save_name) + + elapsed_hyper = time.perf_counter() - start_hyper + print( + "Time for hyperparameter testing: {:.4f} minutes".format( + elapsed_hyper / 60)) + print( + "Time for hyperparameter testing: {:.4f} hours".format( + elapsed_hyper / 60 / 60)) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py index 4a9ac15143..d33e7ad763 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py @@ -160,15 +160,35 @@ def extrapolate_ground_truth_data(data_dir, num_days=360): print(f"Ground truth data saved to {output_path}") +def generate_bounds(data_dir): + path = data_dir + "/Germany/pydata/ground_truth_all_nodes.pickle" + with open(path, 'rb') as f: + ground_truth_all_nodes_np = pd.read_pickle(f) + + # Calculate the bounds + lower_bound = np.min(ground_truth_all_nodes_np, axis=1) + upper_bound = np.max(ground_truth_all_nodes_np, axis=1) + path_upper_bound = data_dir + "/Germany/pydata/ground_truth_upper_bound.pickle" + path_lower_bound = data_dir + "/Germany/pydata/ground_truth_lower_bound.pickle" + with open(path_upper_bound, 'wb') as f: + pd.to_pickle(upper_bound, f) + with open(path_lower_bound, 'wb') as f: + pd.to_pickle(lower_bound, f) + + print(f"Upper bound saved to {path_upper_bound}") + print(f"Lower bound saved to {path_lower_bound}") + + def main(): cwd = os.getcwd() - num_days = 360 + num_days = 180 data_dir = os.path.join(cwd, "data") extrapolate_ground_truth_data( data_dir, num_days ) + generate_bounds(data_dir) if __name__ == "__main__": diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py new file mode 100644 index 0000000000..2563de94e6 --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py @@ -0,0 +1,87 @@ +import os +import pickle + +import pandas as pd +import numpy as np +from sklearn.preprocessing import FunctionTransformer +import tensorflow as tf +from tensorflow.keras.layers import Dense +from tensorflow.keras.losses import MeanAbsolutePercentageError +from tensorflow.keras.metrics import mean_absolute_percentage_error +from tensorflow.keras.models import Model +import tensorflow.keras.initializers as initializers + +from tensorflow.keras.optimizers import Adam, Nadam, RMSprop, SGD, Adagrad +import spektral.layers + +from sklearn.model_selection import KFold + +# from spektral.data import Dataset, DisjointLoader, Graph, Loader, BatchLoader, MixedLoader +# from spektral.layers import GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv, AGNNConv, APPNPConv, CrystalConv, GATConv, GINConv, XENetConv, GCNConv, GCSConv +# from spektral.transforms.normalize_adj import NormalizeAdj +# from spektral.utils.convolution import gcn_filter, normalized_laplacian, rescale_laplacian, normalized_adjacency + + +# Function to generate a model class dynamically + +def generate_model_class(name, layer_types, num_repeat, num_output, transform=None): + def __init__(self): + super(type(self), self).__init__() + self.layer_seq = [] + for i, layer_type in enumerate(layer_types): + for _ in range(num_repeat[i]): + layer = layer_type() if callable(layer_type) else layer_type + self.layer_seq.append(layer) + self.output_layer = tf.keras.layers.Dense( + num_output, activation="relu") + + def call(self, inputs): + # Input consists of a tuple (x, a) where x is the node features and a is the adjacency matrix + x, a = inputs + + # Apply transformation on adjacency matrix if provided + if transform is not None: + a = transform(a) + + # Pass through the layers + for layer in self.layer_seq: + if type(layer).__module__.startswith("spektral.layers"): + x = layer([x, a]) # Pass both `x` and `a` to the layer + else: + x = layer(x) # Pass only `x` to the layer + + output = self.output_layer(x) + return output + + # Define the methods + class_dict = { + '__init__': __init__, + 'call': call + } + return type(name, (tf.keras.Model,), class_dict) + + +if __name__ == "__main__": + layer_types = [ + # Dense layer (only uses x) + lambda: tf.keras.layers.Dense(10, activation="relu"), + # Dense layer (only uses x) + lambda: tf.keras.layers.Dense(20, activation="relu"), + lambda: tf.keras.layers.Dense(30, activation="relu") + ] + num_repeat = [2, 3, 1] + + Karlheinz = generate_model_class( + "Rosmarie", layer_types, num_repeat, num_output=2) + model = Karlheinz() + + # Example inputs + x = tf.random.normal([10, 20]) # Node features + a = tf.random.normal([10, 10]) # Adjacency matrix + labels = tf.random.normal([10, 2]) # Example labels + + model.compile(optimizer="adam", loss="mse") + model.fit([x, a], labels, epochs=5) + + # Print trainable variables + print(model.trainable_variables) From 306c7526ef484aa2a8bd70d4b09be6a07780e61d Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Wed, 3 Sep 2025 13:46:23 +0200 Subject: [PATCH 06/41] [ci skip] Include different network architectures and introduce grid search procedure. --- .../surrogatemodel/GNN/evaluate_and_train.py | 93 ++++++++------ .../memilio/surrogatemodel/GNN/grid_search.py | 119 ++++++++++++++++++ .../GNN/network_architectures.py | 66 ++++++++-- 3 files changed, 230 insertions(+), 48 deletions(-) create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index 4759ad194f..af32cf7dbc 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -19,7 +19,8 @@ from spektral.transforms.normalize_adj import NormalizeAdj from spektral.data import Dataset, DisjointLoader, Graph, Loader, BatchLoader, MixedLoader -from spektral.layers import GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv +from spektral.layers import (GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv, AGNNConv, + APPNPConv, CrystalConv, GATConv, GINConv, XENetConv, GCNConv, GCSConv) from spektral.utils.convolution import normalized_laplacian, rescale_laplacian @@ -111,7 +112,7 @@ def evaluate(loader, model, loss_fn, retransform=False): return np.average(output[:, :-1], 0, weights=output[:, -1]) -def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_patience, save_name): +def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_patience, save_results=False, save_name=""): n = len(data) n_train, n_valid, n_test = calc_split_index( n, split_train=0.7, split_valid=0.2, split_test=0.1) @@ -215,44 +216,56 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ print(f"Time for training: {elapsed:.4f} seconds") print(f"Time for training: {elapsed/60:.4f} minutes") - -# save df - df.loc[len(df.index)] = [ # layer_name, number_of_layer, channels, - np.mean(train_losses), - np.mean(val_losses), - np.mean(test_scores), - np.mean(test_scores_r), - (elapsed / 60), - [losses_history_all], - [val_losses_history_all]] - print(df) - - # Ensure that save_name has the .pickle extension - if not save_name.endswith('.pickle'): - save_name += '.pickle' - - # Save best weights as pickle - path = os.path.dirname(os.path.realpath(__file__)) - file_path_w = os.path.join(path, 'saved_weights') - # Ensure the directory exists - if not os.path.isdir(file_path_w): - os.mkdir(file_path_w) - - # Construct the full file path by joining the directory with save_name - file_path_w = os.path.join(file_path_w, save_name) - - # Save the weights to the file - with open(file_path_w, 'wb') as f: - pickle.dump(best_weights, f) - - file_path_df = os.path.join( - os.path.dirname( - os.path.realpath(os.path.dirname(os.path.realpath(path)))), - 'model_evaluations_paper') - if not os.path.isdir(file_path_df): - os.mkdir(file_path_df) - file_path_df = file_path_df+save_name.replace('.pickle', '.csv') - df.to_csv(file_path_df) + if save_results: + # save df + df.loc[len(df.index)] = [ # layer_name, number_of_layer, channels, + np.mean(train_losses), + np.mean(val_losses), + np.mean(test_scores), + np.mean(test_scores_r), + (elapsed / 60), + [losses_history_all], + [val_losses_history_all]] + print(df) + + # Ensure that save_name has the .pickle extension + if not save_name.endswith('.pickle'): + save_name += '.pickle' + + # Save best weights as pickle + path = os.path.dirname(os.path.realpath(__file__)) + file_path_w = os.path.join(path, 'saved_weights') + # Ensure the directory exists + if not os.path.isdir(file_path_w): + os.mkdir(file_path_w) + + # Construct the full file path by joining the directory with save_name + file_path_w = os.path.join(file_path_w, save_name) + + # Save the weights to the file + with open(file_path_w, 'wb') as f: + pickle.dump(best_weights, f) + + file_path_df = os.path.join( + os.path.dirname( + os.path.realpath(os.path.dirname(os.path.realpath(path)))), + 'model_evaluations_paper') + if not os.path.isdir(file_path_df): + os.mkdir(file_path_df) + file_path_df = file_path_df+save_name.replace('.pickle', '.csv') + df.to_csv(file_path_df) + + else: + return { + "model": save_name, + "mean_train_loss": np.mean(train_losses), + "mean_val_loss": np.mean(val_losses), + "mean_test_loss": np.mean(test_scores), + "mean_test_loss_orig": np.mean(test_scores_r), + "training_time": elapsed/60, + "train_losses": [losses_history_all], + "val_losses": [val_losses_history_all] + } if __name__ == "__main__": diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py new file mode 100644 index 0000000000..993e3c812e --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py @@ -0,0 +1,119 @@ +import os +import pickle + +import pandas as pd +import numpy as np +from sklearn.preprocessing import FunctionTransformer +import tensorflow as tf +from tensorflow.keras.layers import Dense +from tensorflow.keras.losses import MeanAbsolutePercentageError +from tensorflow.keras.metrics import mean_absolute_percentage_error +from tensorflow.keras.models import Model +import tensorflow.keras.initializers as initializers + +from tensorflow.keras.optimizers import Adam, Nadam, RMSprop, SGD, Adagrad +import spektral.layers as spektral_layers +import spektral.utils.convolution as spektral_convolution +from spektral.data import MixedLoader + +from sklearn.model_selection import KFold + +from memilio.surrogatemodel.GNN.network_architectures import get_model +from memilio.surrogatemodel.GNN.evaluate_and_train import ( + train_and_evaluate, create_dataset) + +layer_types = [ + "ARMAConv", + "GCSConv", + "GATConv", + "GCNConv", + "APPNPConv" +] + +numbers_layers = [2, 3, 4, 5, 6, 7] +numbers_channels = [2, 3, 4, 5, 6, 7] +activation_functions = ["elu", "relu", "tanh", "sigmoid"] + +output_dim = 1440 + +batch_size = 32 +loss_function = MeanAbsolutePercentageError() +optimizer = Adam() +es_patience = 30 +max_epochs = 100 + +model_parameters = [(layer, num_layers, num_channels, activation) for layer in layer_types + for num_layers in numbers_layers + for num_channels in numbers_channels + for activation in activation_functions] + + +def perform_grid_search(model_parameters, data): + """Perform grid search over the specified model parameters. + Trains and evaluates models with different configurations and stores the results in a CSV file. + """ + # Create a DataFrame to store the results + df_results = pd.DataFrame(columns=['model', 'optimizer', 'number_of_hidden_layers', 'number_of_channels', 'activation', + 'mean_train_loss', + 'mean_validation_loss', 'training_time', 'train_losses', 'val_losses']) + + for param in model_parameters: + layer_type, num_layers, num_channels, activation = param + print( + f"Training model with {layer_type}, {num_layers} layers, {num_channels} channels, activation: {activation}") + + # Create a model instance + model = get_model( + layer_type=layer_type, + num_layers=num_layers, + num_channels=num_channels, + activation=activation, + num_output=output_dim + ) + optimizer = Adam() + + loader = MixedLoader(data, batch_size=batch_size, epochs=1) + inputs, _ = loader.__next__() + model(inputs) # Build the model by calling it on a batch of data + # Initialize optimizer variables + optimizer.build(model.trainable_variables) + model.compile( + optimizer=optimizer, + loss=MeanAbsolutePercentageError(), + metrics=[mean_absolute_percentage_error] + ) + + results = train_and_evaluate( + data, batch_size, epochs=max_epochs, model=model, + loss_fn=loss_function, optimizer=optimizer, + es_patience=es_patience, save_name="") + + df_results.loc[len(df_results.index)] = [ + layer_type, optimizer.__class__.__name__, num_layers, num_channels, activation, + results["mean_train_loss"], + results["mean_val_loss"], + results["training_time"], + results["train_losses"], + results["val_losses"] + ] + + path = os.path.dirname(os.path.realpath(__file__)) + + file_path_df = os.path.join(os.path.join( + os.path.dirname( + os.path.realpath(os.path.dirname(os.path.realpath(path)))), 'saves')) + if not os.path.isdir(file_path_df): + os.mkdir(file_path_df) + + df_results.to_csv(os.path.join( + file_path_df, 'grid_search_results.csv')) + print(f"Saved intermediate results to {file_path_df}") + tf.keras.backend.clear_session() + + +if __name__ == "__main__": + # Generate the Dataset + path_cases = "/localdata1/hege_mn/memilio/saves/GNN_data_30days_3dampings_classic5.pickle" + path_mobility = '/localdata1/hege_mn/memilio/data/Germany/mobility' + data = create_dataset(path_cases, path_mobility) + perform_grid_search(model_parameters, data) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py index 2563de94e6..6ab089f964 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py @@ -12,7 +12,9 @@ import tensorflow.keras.initializers as initializers from tensorflow.keras.optimizers import Adam, Nadam, RMSprop, SGD, Adagrad -import spektral.layers +import spektral.layers as spektral_layers +import spektral.utils.convolution as spektral_convolution + from sklearn.model_selection import KFold @@ -40,8 +42,9 @@ def call(self, inputs): x, a = inputs # Apply transformation on adjacency matrix if provided - if transform is not None: - a = transform(a) + if not tf.is_symbolic_tensor(a): + if transform is not None: + a = transform(a) # Pass through the layers for layer in self.layer_seq: @@ -61,6 +64,51 @@ def call(self, inputs): return type(name, (tf.keras.Model,), class_dict) +def get_model(layer_type, num_layers, num_channels, activation, num_output=1): + if layer_type == "ARMAConv": + layer_name = spektral_layers.ARMAConv + + def transform(a): + a = np.array(a) + a = spektral_convolution.rescale_laplacian( + spektral_convolution.normalized_laplacian(a)) + return tf.convert_to_tensor(a, dtype=tf.float32) + elif layer_type == "GCSConv": + layer_name = spektral_layers.GCSConv + + def transform(a): + a = a.numpy() + a = spektral_convolution.normalized_adjacency(a) + return tf.convert_to_tensor(a, dtype=tf.float32) + elif layer_type == "GATConv": + layer_name = spektral_layers.GATConv + + def transform(a): + a = a.numpy() + a = spektral_convolution.normalized_adjacency(a) + return tf.convert_to_tensor(a, dtype=tf.float32) + elif layer_type == "GCNConv": + layer_name = spektral_layers.GCNConv + + def transform(a): + a = a.numpy() + a = spektral_convolution.gcn_filter(a) + return tf.convert_to_tensor(a, dtype=tf.float32) + elif layer_type == "APPNPConv": + layer_name = spektral_layers.APPNPConv + + def transform(a): + a = a.numpy() + a = spektral_convolution.gcn_filter(a) + return tf.convert_to_tensor(a, dtype=tf.float32) + + layer_types = [lambda: layer_name(num_channels, activation=activation)] + num_repeat = [num_layers] + model_class = generate_model_class( + "CustomModel", layer_types, num_repeat, num_output, transform=transform) + return model_class() + + if __name__ == "__main__": layer_types = [ # Dense layer (only uses x) @@ -71,9 +119,9 @@ def call(self, inputs): ] num_repeat = [2, 3, 1] - Karlheinz = generate_model_class( - "Rosmarie", layer_types, num_repeat, num_output=2) - model = Karlheinz() + testclass = generate_model_class( + "testclass", layer_types, num_repeat, num_output=2) + model = testclass() # Example inputs x = tf.random.normal([10, 20]) # Node features @@ -83,5 +131,7 @@ def call(self, inputs): model.compile(optimizer="adam", loss="mse") model.fit([x, a], labels, epochs=5) - # Print trainable variables - print(model.trainable_variables) + model2 = get_model("ARMAConv", 3, 16, "relu", num_output=2) + + model2.compile(optimizer="adam", loss="mse") + model2.fit([x, a], labels, epochs=5) From bf6c5f5be7af4f81c1c6d58af1ecd7e74cfdc066 Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 8 Sep 2025 13:56:44 +0200 Subject: [PATCH 07/41] [ci skip] Improve model creation and implement grid search functionality --- .../memilio/surrogatemodel/GNN/GNN_utils.py | 19 + .../surrogatemodel/GNN/evaluate_and_train.py | 94 ++++- .../memilio/surrogatemodel/GNN/grid_search.py | 63 ++- .../GNN/network_architectures.py | 76 +++- .../surrogatemodel/GNN/save_ground_truth.py | 54 +++ .../test_surrogatemodel_GNN.py | 378 ++++++++++++++++++ 6 files changed, 622 insertions(+), 62 deletions(-) create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/save_ground_truth.py create mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py index 53c30e79db..1ede0363cf 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -1,3 +1,22 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Agatha Schmidt, Henrik Zunker +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + import numpy as np import pandas as pd import os diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index af32cf7dbc..36f8e0b056 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -1,55 +1,76 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Agatha Schmidt, Henrik Zunker, Manuel Heger +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + import os import pickle import spektral import time import pandas as pd import numpy as np -from sklearn.preprocessing import FunctionTransformer from tensorflow.keras.optimizers import Adam from tensorflow.keras.losses import MeanAbsolutePercentageError -from tensorflow.keras.models import Model import tensorflow.keras.initializers as initializers -import copy import tensorflow as tf import memilio.surrogatemodel.GNN.network_architectures as network_architectures from memilio.surrogatemodel.utils.helper_functions import (calc_split_index) -from spektral.transforms.normalize_adj import NormalizeAdj -from spektral.data import Dataset, DisjointLoader, Graph, Loader, BatchLoader, MixedLoader +from spektral.data import MixedLoader from spektral.layers import (GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv, AGNNConv, APPNPConv, CrystalConv, GATConv, GINConv, XENetConv, GCNConv, GCSConv) from spektral.utils.convolution import normalized_laplacian, rescale_laplacian -def create_dataset(path_cases, path_mobility): +def create_dataset(path_cases, path_mobility, number_of_nodes=400): """ Create a dataset from the given file path. + :param path_cases: Path to the pickle file containing case data. + :param path_mobility: Path to the directory containing mobility data. + .param number_of_nodes: Number of nodes in the graph. + :return: A Spektral Dataset object containing the processed data. """ + + # Load data from pickle file file = open( path_cases, 'rb') data = pickle.load(file) - - number_of_nodes = 400 - + # Extract inputs and labels from the loaded data inputs = data['inputs'] labels = data['labels'] len_dataset = len(inputs) - + # Calculate the flattened shape of inputs and labels shape_input_flat = np.asarray( inputs).shape[1]*np.asarray(inputs).shape[2] shape_labels_flat = np.asarray( labels).shape[1]*np.asarray(labels).shape[2] - + # Reshape inputs and labels to the required format new_inputs = np.asarray( inputs.transpose(0, 3, 1, 2)).reshape( len_dataset, number_of_nodes, shape_input_flat) new_labels = np.asarray(labels.transpose(0, 3, 1, 2)).reshape( len_dataset, number_of_nodes, shape_labels_flat) + # Load mobility data and create adjacency matrix commuter_file = open(os.path.join( path_mobility, 'commuter_mobility_2022.txt'), 'rb') commuter_data = pd.read_csv(commuter_file, sep=" ", header=None) @@ -58,32 +79,38 @@ def create_dataset(path_cases, path_mobility): adjacency_matrix = np.asarray(sub_matrix) adjacency_matrix[adjacency_matrix > 0] = 1 # make the adjacency binary + node_features = new_inputs node_labels = new_labels + # Define a custom Dataset class class MyDataset(spektral.data.dataset.Dataset): def read(self): self.a = adjacency_matrix - - # self.a = normalized_adjacency(adjacency_matrix) - # self.a = rescale_laplacian(normalized_laplacian(adjacency_matrix)) - return [spektral.data.Graph(x=x, y=y) for x, y in zip(node_features, node_labels)] super().__init__(**kwargs) - - # data = MyDataset() + # Instantiate the custom dataset data = MyDataset() return data def train_step(inputs, target, loss_fn, model, optimizer): + '''Perform a single training step. + :param inputs: Tuple (x, a) where x is the node features and a is the adjacency matrix. + :param target: Ground truth labels. + :param loss_fn: Loss function to use. + :param model: The GNN model to train. + :param optimizer: Optimizer to use for training. + :return: Loss and accuracy for the training step.''' + # Record operations for automatic differentiation with tf.GradientTape() as tape: predictions = model(inputs, training=True) loss = loss_fn(target, predictions) + \ sum(model.losses) # Add regularization losses + # Compute gradients and update model weights gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) acc = tf.reduce_mean(loss_fn(target, predictions)) @@ -91,6 +118,12 @@ def train_step(inputs, target, loss_fn, model, optimizer): def evaluate(loader, model, loss_fn, retransform=False): + '''Evaluate the model on a validation or test set. + :param loader: Data loader for the evaluation set. + :param model: The GNN model to evaluate. + :param loss_fn: Loss function to use. + :param retransform: Whether to apply inverse transformation to the outputs. + :return: Losses and mean_loss for the evaluation set.''' output = [] step = 0 while step < loader.steps_per_epoch: @@ -107,13 +140,26 @@ def evaluate(loader, model, loss_fn, retransform=False): len(target), # Keep track of batch size ) output.append(outs) + # Aggregate results at the end of the epoch if step == loader.steps_per_epoch: output = np.array(output) return np.average(output[:, :-1], 0, weights=output[:, -1]) -def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_patience, save_results=False, save_name=""): +def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_patience, save_results=False, save_name="model"): + '''Train and evaluate the GNN model. + :param data: The dataset to use for training and evaluation. + :param batch_size: Batch size for training. + :param epochs: Maximum number of epochs to train. + :param model: The GNN model to train and evaluate. + :param loss_fn: Loss function to use. + :param optimizer: Optimizer to use for training. + :param es_patience: Patience for early stopping. + :param save_results: Whether to save the results to a file. + :param save_name: Name to use when saving the results. + :return: A dictionary containing training and evaluation results if save_results is False.''' n = len(data) + # Determine split indices for training, validation, and test sets n_train, n_valid, n_test = calc_split_index( n, split_train=0.7, split_valid=0.2, split_test=0.1) @@ -134,6 +180,7 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ "test_loss_orig", "training_time", "loss_history", "val_loss_history"]) + # Initialize variables to track best scores and histories test_scores = [] test_scores_r = [] train_losses = [] @@ -149,6 +196,9 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ losses_history = [] val_losses_history = [] + ################################################################################ + # Train model + ################################################################################ start = time.perf_counter() step = 0 epoch = 0 @@ -157,7 +207,7 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ loss, acc = train_step(*batch, loss_fn, model, optimizer) results.append((loss, acc)) - # Compute validation loss and accuracy + # Compute validation loss and accuracy at the end of each epoch if step == loader_tr.steps_per_epoch: step = 0 epoch += 1 @@ -186,7 +236,8 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ ################################################################################ # Evaluate model ################################################################################ - model.set_weights(best_weights) # Load best model + # Load best model weights before evaluation on test set + model.set_weights(best_weights) test_loss, test_acc = evaluate(loader_test, model, loss_fn) test_loss_r, test_acc_r = evaluate( loader_test, model, loss_fn, retransform=True) @@ -216,6 +267,9 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ print(f"Time for training: {elapsed:.4f} seconds") print(f"Time for training: {elapsed/60:.4f} minutes") + ################################################################################ + # Save results + ################################################################################ if save_results: # save df df.loc[len(df.index)] = [ # layer_name, number_of_layer, channels, diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py index 993e3c812e..2b2b5850de 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py @@ -1,27 +1,37 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Agatha Schmidt, Henrik Zunker, Manuel Heger +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + import os -import pickle import pandas as pd -import numpy as np -from sklearn.preprocessing import FunctionTransformer import tensorflow as tf -from tensorflow.keras.layers import Dense from tensorflow.keras.losses import MeanAbsolutePercentageError from tensorflow.keras.metrics import mean_absolute_percentage_error -from tensorflow.keras.models import Model -import tensorflow.keras.initializers as initializers -from tensorflow.keras.optimizers import Adam, Nadam, RMSprop, SGD, Adagrad -import spektral.layers as spektral_layers -import spektral.utils.convolution as spektral_convolution +from tensorflow.keras.optimizers import Adam from spektral.data import MixedLoader - -from sklearn.model_selection import KFold - from memilio.surrogatemodel.GNN.network_architectures import get_model from memilio.surrogatemodel.GNN.evaluate_and_train import ( train_and_evaluate, create_dataset) +# Define the parameter grid for grid search layer_types = [ "ARMAConv", "GCSConv", @@ -33,25 +43,34 @@ numbers_layers = [2, 3, 4, 5, 6, 7] numbers_channels = [2, 3, 4, 5, 6, 7] activation_functions = ["elu", "relu", "tanh", "sigmoid"] +model_parameters = [(layer, num_layers, num_channels, activation) for layer in layer_types + for num_layers in numbers_layers + for num_channels in numbers_channels + for activation in activation_functions] -output_dim = 1440 - +# Fix training parameters batch_size = 32 loss_function = MeanAbsolutePercentageError() optimizer = Adam() es_patience = 30 max_epochs = 100 -model_parameters = [(layer, num_layers, num_channels, activation) for layer in layer_types - for num_layers in numbers_layers - for num_channels in numbers_channels - for activation in activation_functions] +training_parameters = [batch_size, loss_function, + optimizer, es_patience, max_epochs] -def perform_grid_search(model_parameters, data): +def perform_grid_search(model_parameters, training_parameters, data): """Perform grid search over the specified model parameters. Trains and evaluates models with different configurations and stores the results in a CSV file. + Results are saved in 'grid_search_results.csv' in the 'saves' directory. + :param model_parameters: List of tuples containing model parameters (layer_type, num_layers, num_channels, activation). + :param training_parameters: List containing training parameters (batch_size, loss_function, optimizer, es_patience, max_epochs). + :param data: The dataset to be used for training and evaluation. + :return: None """ + # Setting hyprparameter + output_dim = data[0].y.shape[-1] + batch_size, loss_function, optimizer, es_patience, max_epochs = training_parameters # Create a DataFrame to store the results df_results = pd.DataFrame(columns=['model', 'optimizer', 'number_of_hidden_layers', 'number_of_channels', 'activation', 'mean_train_loss', @@ -75,6 +94,7 @@ def perform_grid_search(model_parameters, data): loader = MixedLoader(data, batch_size=batch_size, epochs=1) inputs, _ = loader.__next__() model(inputs) # Build the model by calling it on a batch of data + # Initialize optimizer variables optimizer.build(model.trainable_variables) model.compile( @@ -96,7 +116,7 @@ def perform_grid_search(model_parameters, data): results["train_losses"], results["val_losses"] ] - + # Save intermediate results to avoid data loss path = os.path.dirname(os.path.realpath(__file__)) file_path_df = os.path.join(os.path.join( @@ -106,8 +126,9 @@ def perform_grid_search(model_parameters, data): os.mkdir(file_path_df) df_results.to_csv(os.path.join( - file_path_df, 'grid_search_results.csv')) + file_path_df, 'grid_search_results.csv'), index=False) print(f"Saved intermediate results to {file_path_df}") + # Clear session to free memory after each iteration tf.keras.backend.clear_session() diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py index 6ab089f964..0fd0da563f 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py @@ -1,32 +1,41 @@ -import os -import pickle - -import pandas as pd +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Agatha Schmidt, Henrik Zunker, Manuel Heger +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# import numpy as np -from sklearn.preprocessing import FunctionTransformer import tensorflow as tf -from tensorflow.keras.layers import Dense -from tensorflow.keras.losses import MeanAbsolutePercentageError -from tensorflow.keras.metrics import mean_absolute_percentage_error -from tensorflow.keras.models import Model -import tensorflow.keras.initializers as initializers -from tensorflow.keras.optimizers import Adam, Nadam, RMSprop, SGD, Adagrad import spektral.layers as spektral_layers import spektral.utils.convolution as spektral_convolution -from sklearn.model_selection import KFold - -# from spektral.data import Dataset, DisjointLoader, Graph, Loader, BatchLoader, MixedLoader -# from spektral.layers import GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv, AGNNConv, APPNPConv, CrystalConv, GATConv, GINConv, XENetConv, GCNConv, GCSConv -# from spektral.transforms.normalize_adj import NormalizeAdj -# from spektral.utils.convolution import gcn_filter, normalized_laplacian, rescale_laplacian, normalized_adjacency - - # Function to generate a model class dynamically def generate_model_class(name, layer_types, num_repeat, num_output, transform=None): + '''Generates a custom Keras model class with specified layers. + + :param name: Name of the generated class. + :param layer_types: List of layer types (classes or callables) to include in the model. + :param num_repeat: List of integers specifying how many times to repeat each layer type. + :param num_output: Number of output units in the final layer. + :param transform: Optional function to transform the adjacency matrix before passing it to layers. + :return: A dynamically created Keras model class.''' + def __init__(self): super(type(self), self).__init__() self.layer_seq = [] @@ -34,6 +43,7 @@ def __init__(self): for _ in range(num_repeat[i]): layer = layer_type() if callable(layer_type) else layer_type self.layer_seq.append(layer) + # Final output layer self.output_layer = tf.keras.layers.Dense( num_output, activation="relu") @@ -65,6 +75,29 @@ def call(self, inputs): def get_model(layer_type, num_layers, num_channels, activation, num_output=1): + """Generates a GNN model based on the specified parameters. + :param layer_type: Name of GNN layer to use (possible 'ARMAConv', 'GCSConv', 'GATConv', + 'GCNConv', 'APPNPConv'), provided as a string. + :param num_layers: Number of hidden layers in the model. + :param num_channels: Number of channels (units) in each hidden layer. + :param activation: Activation function to use in the hidden layers (e.g., 'relu', 'elu', 'tanh', 'sigmoid'). + :param num_output: Number of output units in the final layer. + :return: A Keras model instance with the specified architecture. + """ + if layer_type not in ["ARMAConv", "GCSConv", "GATConv", "GCNConv", "APPNPConv"]: + raise ValueError( + f"Unsupported layer_type: {layer_type}. Supported types are 'ARMAConv', 'GCSConv', 'GATConv', 'GCNConv', 'APPNPConv'.") + if num_layers < 1: + raise ValueError("num_layers must be at least 1.") + if num_channels < 1: + raise ValueError("num_channels must be at least 1.") + if not isinstance(activation, str): + raise ValueError( + "activation must be a string representing the activation function.") + if num_output < 1: + raise ValueError("num_output must be at least 1.") + + # Define the layer based on the specified type if layer_type == "ARMAConv": layer_name = spektral_layers.ARMAConv @@ -101,8 +134,9 @@ def transform(a): a = a.numpy() a = spektral_convolution.gcn_filter(a) return tf.convert_to_tensor(a, dtype=tf.float32) - - layer_types = [lambda: layer_name(num_channels, activation=activation)] + # Generate the input for generate_model_class + layer_types = [lambda: layer_name( + num_channels, activation=activation, kernel_initializer=tf.keras.initializers.GlorotUniform())] num_repeat = [num_layers] model_class = generate_model_class( "CustomModel", layer_types, num_repeat, num_output, transform=transform) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/save_ground_truth.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/save_ground_truth.py new file mode 100644 index 0000000000..3b8310bf8f --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/save_ground_truth.py @@ -0,0 +1,54 @@ +import os +import numpy as np +import pandas as pd +import json +import matplotlib.pyplot as plt +import h5py + + +def read_case_data(): + df = pd.read_hdf( + "/localdata1/hege_mn/memilio/data/Germany/pydata/Results_rki.h5") + return df + + +def read_results_h5(path, group_key='Group1'): + with h5py.File(path, 'r') as f: + keys = list(f.keys()) + res = {} + for i, key in enumerate(keys): + group = f[key] + total = group[group_key][()] + # remove confirmed compartments + sum_inf_no_symp = np.sum(total[:, [2, 3]], axis=1) + sum_inf_symp = np.sum(total[:, [4, 5]], axis=1) + total[:, 2] = sum_inf_no_symp + total[:, 4] = sum_inf_symp + total = np.delete(total, [3, 5], axis=1) + res[key] = total + return res + + +def main(): + path_rki_h5 = "/localdata1/hege_mn/memilio/data/Germany/pydata/Results_rki.h5" + num_age_groups = 6 + all_age_data_list = [] + group_keys = ['Group1', 'Group2', 'Group3', 'Group4', 'Group5', 'Group6'] + for age in range(num_age_groups): + age_data_dict = read_results_h5(path_rki_h5, group_keys[age]) + age_data_np = np.array(list(age_data_dict.values())) + all_age_data_list.append(age_data_np) + # Combine age groups to get shape (num_nodes, timesteps, num_features=48) + all_age_data = np.stack(all_age_data_list, axis=-1) + print(f"Shape of all_age_data: {all_age_data.shape}") + num_nodes, timesteps, _, _ = all_age_data.shape + + # Save the ground truth data + output_path = "/localdata1/hege_mn/memilio/data/Germany/pydata/ground_truth_all_nodes.pickle" + with open(output_path, 'wb') as f: + pd.to_pickle(all_age_data, f) + print(f"Ground truth data saved to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py new file mode 100644 index 0000000000..8c7a8b62b1 --- /dev/null +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -0,0 +1,378 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Manuel Heger, Henrik Zunker +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + +from pyfakefs import fake_filesystem_unittest +import memilio.surrogatemodel.GNN.network_architectures as gnn_arch + +from unittest.mock import patch +import os +import unittest +import pickle +import numpy as np +import pandas as pd +import tensorflow as tf +import logging + +# suppress all autograph warnings from tensorflow + +logging.getLogger("tensorflow").setLevel(logging.ERROR) + + +class TestSurrogatemodelGNN(fake_filesystem_unittest.TestCase): + path = "/home/" + + def setUp(self): + self.setUpPyfakefs() + + def test_generate_model_class(self): + from memilio.surrogatemodel.GNN.network_architectures import generate_model_class + + # Test parameters + layer_types = [ + lambda: tf.keras.layers.Dense(10, activation="relu"), + ] + num_layers = [3] + num_output = 2 + + # Generate the model class + ModelClass = generate_model_class( + "TestModel", layer_types, num_layers, num_output) + # Check if the generated class is a subclass of tf.keras.Model + self.assertTrue(issubclass(ModelClass, tf.keras.Model)) + + # Instantiate the model + model = ModelClass() + + # Check if the model has the expected number of layers + expected_num_layers = num_layers[0] + 1 # +1 for the output layer + self.assertEqual(len(model.layers), expected_num_layers) + self.assertIsInstance(model.layers[-1], tf.keras.layers.Dense) + self.assertEqual(model.layers[-1].units, num_output) + self.assertEqual( + model.layers[-1].activation.__name__, "relu") + + def test_get_model(self): + from memilio.surrogatemodel.GNN.network_architectures import get_model + + # Test parameters + layer_type = "GCNConv" + num_layers = 2 + num_channels = 16 + activation = "relu" + num_output = 3 + + # Generate the model + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + + # Check if the model is an instance of tf.keras.Model + self.assertIsInstance(model, tf.keras.Model) + + # Check if the model has the expected number of layers + expected_num_layers = num_layers + 1 # +1 for the output layer + self.assertEqual(len(model.layers), expected_num_layers) + + layer_type = "MonvConv" + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual(str( + error.exception), "Unsupported layer_type: MonvConv. " + "Supported types are 'ARMAConv', 'GCSConv', 'GATConv', 'GCNConv', 'APPNPConv'.") + layer_type = "GATConv" + num_layers = 0 + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual(str( + error.exception), "num_layers must be at least 1.") + + num_layers = 2 + num_output = 0 + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual(str( + error.exception), "num_output must be at least 1.") + + num_output = 2 + num_channels = 0 + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual(str( + error.exception), "num_channels must be at least 1.") + + num_channels = 16 + activation = 5 + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual(str( + error.exception), "activation must be a string representing the activation function.") + + def test_create_dataset(self): + from memilio.surrogatemodel.GNN.evaluate_and_train import create_dataset + + # Create dummy data in the fake filesystem for testing + num_samples = 10 + num_nodes = 5 + num_node_features = 3 + output_dim = 4 + X = np.random.rand(num_samples, 1, + num_node_features, num_nodes).astype(np.float32) + A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) + y = np.random.rand(num_samples, 1, + output_dim, num_nodes).astype(np.float32) + data = {"inputs": X, "labels": y} + path_cases_dir = os.path.join(self.path, "cases") + self.fs.create_dir(path_cases_dir) + path_cases = os.path.join(path_cases_dir, "cases.pickle") + with open(path_cases, 'wb') as f: + pickle.dump(data, f) + path_mobility = os.path.join(self.path, "mobility") + mobility_file = os.path.join( + path_mobility, "commuter_mobility_2022.txt") + self.fs.create_file(mobility_file) + with open(mobility_file, 'w') as f: + np.savetxt(f, A, delimiter=" ") + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + self.assertEqual(len(dataset), num_samples) + for graph in dataset: + self.assertEqual(graph.x.shape, (num_nodes, num_node_features)) + self.assertEqual(dataset.a.shape, (num_nodes, num_nodes)) + self.assertEqual(graph.y.shape, (num_nodes, output_dim)) + # Clean up + self.fs.remove_object(path_cases) + self.fs.remove_object(os.path.join( + path_mobility, "commuter_mobility_2022.txt")) + self.fs.remove_object(path_mobility) + + def test_train_step(self): + from memilio.surrogatemodel.GNN.evaluate_and_train import train_step + from tensorflow.keras.losses import MeanAbsolutePercentageError + + # Create a simple model for testing + model = tf.keras.Sequential([ + tf.keras.layers.Dense(10, activation='relu'), + tf.keras.layers.Dense(2) + ]) + optimizer = tf.keras.optimizers.Adam() + loss_fn = MeanAbsolutePercentageError() + + # Create dummy data + batch_size = 4 + num_nodes = 5 + num_node_features = 3 + output_dim = 2 + inputs = np.random.rand(batch_size, num_nodes, + num_node_features).astype(np.float32) + y = np.random.rand(batch_size, num_nodes, + output_dim).astype(np.float32) + + # Perform a training step + loss, acc = train_step(inputs, y, loss_fn, model, optimizer) + + # Check if the loss is a scalar tensor + self.assertIsInstance(loss, tf.Tensor) + self.assertEqual(loss.shape, ()) + self.assertIsInstance(acc, tf.Tensor) + self.assertEqual(acc.shape, ()) + self.assertGreaterEqual(acc.numpy(), 0) + self.assertGreaterEqual(loss.numpy(), 0) + + def test_evaluate(self): + from memilio.surrogatemodel.GNN.evaluate_and_train import ( + evaluate, create_dataset, MixedLoader) + from memilio.surrogatemodel.GNN.network_architectures import get_model + from tensorflow.keras.losses import MeanAbsolutePercentageError + + # Create a simple model for testing + model = get_model( + layer_type="GCNConv", num_layers=2, num_channels=16, activation="relu", num_output=4) + loss_fn = MeanAbsolutePercentageError() + + # Create dummy data in the fake filesystem for testing + num_samples = 10 + num_nodes = 5 + num_node_features = 3 + output_dim = 4 + X = np.random.rand(num_samples, 1, + num_node_features, num_nodes).astype(np.float32) + A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) + y = np.random.rand(num_samples, 1, + output_dim, num_nodes).astype(np.float32) + data = {"inputs": X, "labels": y} + path_cases_dir = os.path.join(self.path, "cases") + self.fs.create_dir(path_cases_dir) + path_cases = os.path.join(path_cases_dir, "cases.pickle") + with open(path_cases, 'wb') as f: + pickle.dump(data, f) + path_mobility = os.path.join(self.path, "mobility") + mobility_file = os.path.join( + path_mobility, "commuter_mobility_2022.txt") + self.fs.create_file(mobility_file) + with open(mobility_file, 'w') as f: + np.savetxt(f, A, delimiter=" ") + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + loader = MixedLoader(dataset, batch_size=2, epochs=1) + res = evaluate(loader, model, loss_fn) + + self.assertEqual(len(res), 2) + self.assertGreaterEqual(res[0], 0) + self.assertGreaterEqual(res[1], 0) + + # Test with retransformation + loader = MixedLoader(dataset, batch_size=2, epochs=1) + res = evaluate(loader, model, loss_fn, True) + + self.assertEqual(len(res), 2) + self.assertGreaterEqual(res[0], 0) + self.assertGreaterEqual(res[1], 0) + # Clean up + self.fs.remove_object(path_cases) + self.fs.remove_object(os.path.join( + path_mobility, "commuter_mobility_2022.txt")) + self.fs.remove_object(path_mobility) + + def test_train_and_evaluate(self): + from memilio.surrogatemodel.GNN.evaluate_and_train import ( + train_and_evaluate, create_dataset, MixedLoader) + from memilio.surrogatemodel.GNN.network_architectures import get_model + from tensorflow.keras.losses import MeanAbsolutePercentageError + + number_of_epochs = 2 + # Create a simple model for testing + model = get_model( + layer_type="GCNConv", num_layers=2, num_channels=16, activation="relu", num_output=4) + + # Create dummy data in the fake filesystem for testing + num_samples = 20 + num_nodes = 5 + num_node_features = 3 + output_dim = 4 + X = np.random.rand(num_samples, 1, + num_node_features, num_nodes).astype(np.float32) + A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) + y = np.random.rand(num_samples, 1, + output_dim, num_nodes).astype(np.float32) + data = {"inputs": X, "labels": y} + path_cases_dir = os.path.join(self.path, "cases") + self.fs.create_dir(path_cases_dir) + path_cases = os.path.join(path_cases_dir, "cases.pickle") + with open(path_cases, 'wb') as f: + pickle.dump(data, f) + path_mobility = os.path.join(self.path, "mobility") + mobility_file = os.path.join( + path_mobility, "commuter_mobility_2022.txt") + self.fs.create_file(mobility_file) + with open(mobility_file, 'w') as f: + np.savetxt(f, A, delimiter=" ") + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + + res = train_and_evaluate( + dataset, + batch_size=2, + epochs=number_of_epochs, + model=model, + loss_fn=MeanAbsolutePercentageError(), + optimizer=tf.keras.optimizers.Adam(), + es_patience=100) + + self.assertEqual(len(res["train_losses"][0][0]), number_of_epochs) + + @patch("os.path.realpath", return_value="/home/") + def test_perform_grid_search(self, mock_realpath): + from memilio.surrogatemodel.GNN.evaluate_and_train import ( + create_dataset) + from memilio.surrogatemodel.GNN.grid_search import perform_grid_search + from memilio.surrogatemodel.GNN.network_architectures import get_model + from tensorflow.keras.losses import MeanAbsolutePercentageError + + # Create dummy data in the fake filesystem for testing + num_samples = 20 + num_nodes = 5 + num_node_features = 3 + output_dim = 4 + X = np.random.rand(num_samples, 1, + num_node_features, num_nodes).astype(np.float32) + A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) + y = np.random.rand(num_samples, 1, + output_dim, num_nodes).astype(np.float32) + data = {"inputs": X, "labels": y} + path_cases_dir = os.path.join(self.path, "cases") + self.fs.create_dir(path_cases_dir) + path_cases = os.path.join(path_cases_dir, "cases.pickle") + with open(path_cases, 'wb') as f: + pickle.dump(data, f) + path_mobility = os.path.join(self.path, "mobility") + mobility_file = os.path.join( + path_mobility, "commuter_mobility_2022.txt") + self.fs.create_file(mobility_file) + with open(mobility_file, 'w') as f: + np.savetxt(f, A, delimiter=" ") + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + layers = ["GCNConv", "GATConv"] + num_layers = [1, 2] + activation = "relu" + num_channels = 2 + + model_parameters = [ + (layer, n_layer, num_channels, activation) for layer in layers for n_layer in num_layers + ] + batch_size = 2 + loss_function = MeanAbsolutePercentageError() + optimizer = tf.keras.optimizers.Adam() + es_patience = 5 + max_epochs = 2 + training_parameters = [batch_size, loss_function, + optimizer, es_patience, max_epochs] + # Perform grid search + perform_grid_search(model_parameters, training_parameters, + dataset) + + # Check if the results file is created + results_file = os.path.join( + self.path, "saves", "grid_search_results.csv") + self.assertTrue(os.path.exists(results_file)) + + # Check if the results file has the expected number of rows + df_results = pd.read_csv(results_file) + print(df_results.columns) + self.assertEqual(len(df_results), len(model_parameters)) + self.assertEqual(len(df_results.columns), 10) + # Clean up + self.fs.remove_object(path_cases) + self.fs.remove_object(os.path.join( + path_mobility, "commuter_mobility_2022.txt")) + self.fs.remove_object(path_mobility) + self.fs.remove_object(os.path.join(self.path, "saves")) + + +if __name__ == '__main__': + unittest.main() From 5b24fd9c0be4ff63d62bdf096f7006f9e7695a38 Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Wed, 10 Sep 2025 08:05:54 +0200 Subject: [PATCH 08/41] Add validation checks for dataset and model parameters; enhance test data generation --- .../surrogatemodel/GNN/evaluate_and_train.py | 4 + .../GNN/network_architectures.py | 6 + .../test_surrogatemodel_GNN.py | 166 +++++++++++------- 3 files changed, 112 insertions(+), 64 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index 36f8e0b056..19b9cf490c 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -58,6 +58,10 @@ def create_dataset(path_cases, path_mobility, number_of_nodes=400): labels = data['labels'] len_dataset = len(inputs) + + if len_dataset == 0: # check if dataset is empty + raise ValueError( + "Dataset is empty. Please provide a valid dataset with at least one sample.") # Calculate the flattened shape of inputs and labels shape_input_flat = np.asarray( inputs).shape[1]*np.asarray(inputs).shape[2] diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py index 0fd0da563f..fe1bb033bd 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py @@ -37,6 +37,12 @@ def generate_model_class(name, layer_types, num_repeat, num_output, transform=No :return: A dynamically created Keras model class.''' def __init__(self): + if len(layer_types) != len(num_repeat): + raise ValueError( + "layer_types and num_repeat must have the same length.") + if any(n < 1 for n in num_repeat): + raise ValueError("All values in num_repeat must be at least 1.") + super(type(self), self).__init__() self.layer_seq = [] for i, layer_type in enumerate(layer_types): diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index 8c7a8b62b1..8d1a659cf8 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -38,6 +38,49 @@ class TestSurrogatemodelGNN(fake_filesystem_unittest.TestCase): path = "/home/" + def create_dummy_data(self, num_samples, num_nodes, num_node_features, output_dim): + """ + Create dummy data for testing. + + :param num_samples: Number of samples in the dataset. + :param num_nodes: Number of nodes in each graph. + :param num_node_features: Number of features per node. + :param output_dim: Number of output dimensions per node. + :return: A dictionary containing inputs, adjacency matrix, and labels. + """ + X = np.random.rand(num_samples, 1, num_node_features, + num_nodes).astype(np.float32) + A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) + y = np.random.rand(num_samples, 1, output_dim, + num_nodes).astype(np.float32) + return {"inputs": X, "adjacency": A, "labels": y} + + def setup_fake_filesystem(self, fs, path, data): + """ + Save dummy data to the fake file system. + + :param fs: The fake file system object. + :param path: The base path in the fake file system. + :param data: The dummy data dictionary containing inputs, adjacency, and labels. + :return: Paths to the cases and mobility files. + """ + path_cases_dir = os.path.join(path, "cases") + fs.create_dir(path_cases_dir) + path_cases = os.path.join(path_cases_dir, "cases.pickle") + with open(path_cases, 'wb') as f: + pickle.dump({"inputs": data["inputs"], + "labels": data["labels"]}, f) + + path_mobility = os.path.join(path, "mobility") + mobility_file = os.path.join( + path_mobility, "commuter_mobility_2022.txt") + fs.create_dir(path_mobility) + fs.create_file(mobility_file) + with open(mobility_file, 'w') as f: + np.savetxt(f, data["adjacency"], delimiter=" ") + + return path_cases, path_mobility + def setUp(self): self.setUpPyfakefs() @@ -68,6 +111,43 @@ def test_generate_model_class(self): self.assertEqual( model.layers[-1].activation.__name__, "relu") + # Test with invalid parameters + layer_types = [ + lambda: tf.keras.layers.Dense(10, activation="relu"), + ] + num_layers = [0] + num_output = 2 + with self.assertRaises(ValueError) as error: + ModelClass = generate_model_class( + "TestModel", layer_types, num_layers, num_output) + model = ModelClass() + + self.assertEqual(str( + error.exception), "All values in num_repeat must be at least 1.") + num_layers = [3, 2] + with self.assertRaises(ValueError) as error: + ModelClass = generate_model_class( + "TestModel", layer_types, num_layers, num_output) + model = ModelClass() + self.assertEqual(str( + error.exception), "layer_types and num_repeat must have the same length.") + + # Test with multiple layer types + layer_types = [ + lambda: tf.keras.layers.Dense(10, activation="relu"), + lambda: tf.keras.layers.Dense(20, activation="relu"), + ] + num_repeat = [2, 3] + num_output = 4 + + ModelClass = generate_model_class( + "TestModel", layer_types, num_repeat, num_output) + model = ModelClass() + + # Check the number of layers + self.assertEqual(len(model.layer_seq), sum(num_repeat)) + self.assertEqual(model.output_layer.units, num_output) + def test_get_model(self): from memilio.surrogatemodel.GNN.network_architectures import get_model @@ -141,18 +221,12 @@ def test_create_dataset(self): A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) y = np.random.rand(num_samples, 1, output_dim, num_nodes).astype(np.float32) - data = {"inputs": X, "labels": y} - path_cases_dir = os.path.join(self.path, "cases") - self.fs.create_dir(path_cases_dir) - path_cases = os.path.join(path_cases_dir, "cases.pickle") - with open(path_cases, 'wb') as f: - pickle.dump(data, f) - path_mobility = os.path.join(self.path, "mobility") - mobility_file = os.path.join( - path_mobility, "commuter_mobility_2022.txt") - self.fs.create_file(mobility_file) - with open(mobility_file, 'w') as f: - np.savetxt(f, A, delimiter=" ") + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) + # Create dataset dataset = create_dataset( path_cases, path_mobility, number_of_nodes=num_nodes) @@ -216,23 +290,11 @@ def test_evaluate(self): num_nodes = 5 num_node_features = 3 output_dim = 4 - X = np.random.rand(num_samples, 1, - num_node_features, num_nodes).astype(np.float32) - A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) - y = np.random.rand(num_samples, 1, - output_dim, num_nodes).astype(np.float32) - data = {"inputs": X, "labels": y} - path_cases_dir = os.path.join(self.path, "cases") - self.fs.create_dir(path_cases_dir) - path_cases = os.path.join(path_cases_dir, "cases.pickle") - with open(path_cases, 'wb') as f: - pickle.dump(data, f) - path_mobility = os.path.join(self.path, "mobility") - mobility_file = os.path.join( - path_mobility, "commuter_mobility_2022.txt") - self.fs.create_file(mobility_file) - with open(mobility_file, 'w') as f: - np.savetxt(f, A, delimiter=" ") + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) # Create dataset dataset = create_dataset( path_cases, path_mobility, number_of_nodes=num_nodes) @@ -272,23 +334,11 @@ def test_train_and_evaluate(self): num_nodes = 5 num_node_features = 3 output_dim = 4 - X = np.random.rand(num_samples, 1, - num_node_features, num_nodes).astype(np.float32) - A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) - y = np.random.rand(num_samples, 1, - output_dim, num_nodes).astype(np.float32) - data = {"inputs": X, "labels": y} - path_cases_dir = os.path.join(self.path, "cases") - self.fs.create_dir(path_cases_dir) - path_cases = os.path.join(path_cases_dir, "cases.pickle") - with open(path_cases, 'wb') as f: - pickle.dump(data, f) - path_mobility = os.path.join(self.path, "mobility") - mobility_file = os.path.join( - path_mobility, "commuter_mobility_2022.txt") - self.fs.create_file(mobility_file) - with open(mobility_file, 'w') as f: - np.savetxt(f, A, delimiter=" ") + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) # Create dataset dataset = create_dataset( path_cases, path_mobility, number_of_nodes=num_nodes) @@ -309,7 +359,6 @@ def test_perform_grid_search(self, mock_realpath): from memilio.surrogatemodel.GNN.evaluate_and_train import ( create_dataset) from memilio.surrogatemodel.GNN.grid_search import perform_grid_search - from memilio.surrogatemodel.GNN.network_architectures import get_model from tensorflow.keras.losses import MeanAbsolutePercentageError # Create dummy data in the fake filesystem for testing @@ -317,23 +366,11 @@ def test_perform_grid_search(self, mock_realpath): num_nodes = 5 num_node_features = 3 output_dim = 4 - X = np.random.rand(num_samples, 1, - num_node_features, num_nodes).astype(np.float32) - A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) - y = np.random.rand(num_samples, 1, - output_dim, num_nodes).astype(np.float32) - data = {"inputs": X, "labels": y} - path_cases_dir = os.path.join(self.path, "cases") - self.fs.create_dir(path_cases_dir) - path_cases = os.path.join(path_cases_dir, "cases.pickle") - with open(path_cases, 'wb') as f: - pickle.dump(data, f) - path_mobility = os.path.join(self.path, "mobility") - mobility_file = os.path.join( - path_mobility, "commuter_mobility_2022.txt") - self.fs.create_file(mobility_file) - with open(mobility_file, 'w') as f: - np.savetxt(f, A, delimiter=" ") + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) # Create dataset dataset = create_dataset( path_cases, path_mobility, number_of_nodes=num_nodes) @@ -366,6 +403,7 @@ def test_perform_grid_search(self, mock_realpath): print(df_results.columns) self.assertEqual(len(df_results), len(model_parameters)) self.assertEqual(len(df_results.columns), 10) + # Clean up self.fs.remove_object(path_cases) self.fs.remove_object(os.path.join( From 614082be3ec57672652603dd4545d6b6e4ad23fb Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 22 Sep 2025 07:48:19 +0200 Subject: [PATCH 09/41] Enhance data scaling functionality and validation; add tests for scale_data utility --- .../memilio/surrogatemodel/GNN/GNN_utils.py | 21 ++++- .../surrogatemodel/GNN/data_generation.py | 7 +- .../surrogatemodel/GNN/evaluate_and_train.py | 6 +- .../test_surrogatemodel_GNN.py | 91 +++++++++++++++++-- 4 files changed, 107 insertions(+), 18 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py index 1ede0363cf..93f0c4851d 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -133,21 +133,29 @@ def get_population(): return population -def scale_data(data): +def scale_data(data, transform=True, num_compartments=8): """ Apply a logarithmic transformation on the data. :param data: dictionary, containing entries "inputs" and "labels" + :param transform: Boolean, if True apply the transformation, else return the data reshaped. + :param num_compartments: Number of compartments in the epidemiological model. :returns scaled_inputs: Transformed input data :returns scaled_labels: Transformed output data """ - - num_groups = int(np.asarray(data['inputs']).shape[2] / 8) + if not np.issubdtype(np.asarray(data['inputs']).dtype, np.number): + raise ValueError("Input data must be numeric.") + if not np.issubdtype(np.asarray(data['labels']).dtype, np.number): + raise ValueError("Label data must be numeric.") + num_groups = int(np.asarray(data['inputs']).shape[2] / num_compartments) transformer = FunctionTransformer(np.log1p, validate=True) # Scale inputs inputs = np.asarray( data['inputs']).transpose(2, 0, 1, 3).reshape(num_groups * 8, -1) - scaled_inputs = transformer.transform(inputs) + if transform: + scaled_inputs = transformer.transform(inputs) + else: + scaled_inputs = inputs original_shape_input = np.asarray(data['inputs']).shape # Reverse the reshape @@ -163,7 +171,10 @@ def scale_data(data): # Scale labels labels = np.asarray( data['labels']).transpose(2, 0, 1, 3).reshape(num_groups * 8, -1) - scaled_labels = transformer.transform(labels) + if transform: + scaled_labels = transformer.transform(labels) + else: + scaled_labels = labels original_shape_labels = np.asarray(data['labels']).shape # Reverse the reshape diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index 9261356720..1b3269b0bc 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -404,6 +404,7 @@ def generate_data( :param input_width: number of time steps used for model input. :param label_width: number of time steps (days) used as model output/label. :param save_data: Option to deactivate the save of the dataset. Per default True. + :param transform: Option to deactivate the transformation of the data. Per default True. :param damping_method: String specifying the damping method, that should be used. Possible values "classic", "active", "random". :param max_number_damping: Maximal number of possible dampings. :returns: Dictionary containing inputs, labels, contact_matrix, damping_days and damping_factors @@ -473,8 +474,8 @@ def generate_data( f"For Days = {days}, AVG runtime: {np.mean(times)}s, Median runtime: {np.median(times)}s") if save_data: - if transform: - inputs, labels = scale_data(data) + + inputs, labels = scale_data(data, transform) all_data = {"inputs": inputs, "labels": labels, @@ -510,7 +511,7 @@ def generate_data( data_dir = "/localdata1/hege_mn/memilio/data" input_width = 5 number_of_dampings = 0 - num_runs = 100 + num_runs = 1 label_width_list = [30] random.seed(10) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index 19b9cf490c..9ca05be2dc 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -310,7 +310,7 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ 'model_evaluations_paper') if not os.path.isdir(file_path_df): os.mkdir(file_path_df) - file_path_df = file_path_df+save_name.replace('.pickle', '.csv') + file_path_df = file_path_df+"/"+save_name.replace('.pickle', '.csv') df.to_csv(file_path_df) else: @@ -321,8 +321,8 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ "mean_test_loss": np.mean(test_scores), "mean_test_loss_orig": np.mean(test_scores_r), "training_time": elapsed/60, - "train_losses": [losses_history_all], - "val_losses": [val_losses_history_all] + "train_losses": losses_history_all, + "val_losses": val_losses_history_all } diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index 8d1a659cf8..2a1ba4070f 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -20,6 +20,7 @@ from pyfakefs import fake_filesystem_unittest import memilio.surrogatemodel.GNN.network_architectures as gnn_arch +import memilio.surrogatemodel.GNN.GNN_utils as utils from unittest.mock import patch import os @@ -121,9 +122,9 @@ def test_generate_model_class(self): ModelClass = generate_model_class( "TestModel", layer_types, num_layers, num_output) model = ModelClass() - self.assertEqual(str( error.exception), "All values in num_repeat must be at least 1.") + num_layers = [3, 2] with self.assertRaises(ValueError) as error: ModelClass = generate_model_class( @@ -169,6 +170,8 @@ def test_get_model(self): expected_num_layers = num_layers + 1 # +1 for the output layer self.assertEqual(len(model.layers), expected_num_layers) + # Check handling of invalid parameters + # Test with invalid layer type layer_type = "MonvConv" with self.assertRaises(ValueError) as error: model = get_model(layer_type, num_layers, @@ -176,6 +179,7 @@ def test_get_model(self): self.assertEqual(str( error.exception), "Unsupported layer_type: MonvConv. " "Supported types are 'ARMAConv', 'GCSConv', 'GATConv', 'GCNConv', 'APPNPConv'.") + # Test with invalud num_layers layer_type = "GATConv" num_layers = 0 with self.assertRaises(ValueError) as error: @@ -183,7 +187,7 @@ def test_get_model(self): num_channels, activation, num_output) self.assertEqual(str( error.exception), "num_layers must be at least 1.") - + # Test with invalid num_output num_layers = 2 num_output = 0 with self.assertRaises(ValueError) as error: @@ -191,7 +195,7 @@ def test_get_model(self): num_channels, activation, num_output) self.assertEqual(str( error.exception), "num_output must be at least 1.") - + # Test with invalid num_channels num_output = 2 num_channels = 0 with self.assertRaises(ValueError) as error: @@ -199,7 +203,7 @@ def test_get_model(self): num_channels, activation, num_output) self.assertEqual(str( error.exception), "num_channels must be at least 1.") - + # Test with invalid activation num_channels = 16 activation = 5 with self.assertRaises(ValueError) as error: @@ -300,7 +304,7 @@ def test_evaluate(self): path_cases, path_mobility, number_of_nodes=num_nodes) loader = MixedLoader(dataset, batch_size=2, epochs=1) res = evaluate(loader, model, loss_fn) - + # Check if the result is a tuple of (loss, accuracy) self.assertEqual(len(res), 2) self.assertGreaterEqual(res[0], 0) self.assertGreaterEqual(res[1], 0) @@ -318,7 +322,8 @@ def test_evaluate(self): path_mobility, "commuter_mobility_2022.txt")) self.fs.remove_object(path_mobility) - def test_train_and_evaluate(self): + @patch("os.path.realpath", return_value="/home/") + def test_train_and_evaluate(self, mock_realpath): from memilio.surrogatemodel.GNN.evaluate_and_train import ( train_and_evaluate, create_dataset, MixedLoader) from memilio.surrogatemodel.GNN.network_architectures import get_model @@ -352,7 +357,47 @@ def test_train_and_evaluate(self): optimizer=tf.keras.optimizers.Adam(), es_patience=100) - self.assertEqual(len(res["train_losses"][0][0]), number_of_epochs) + self.assertEqual(len(res["train_losses"][0]), number_of_epochs) + self.assertEqual(len(res["val_losses"][0]), number_of_epochs) + self.assertGreater(res["mean_test_loss"], 0) + + # Testing with saving the results + res = train_and_evaluate( + dataset, + batch_size=2, + epochs=number_of_epochs, + model=model, + loss_fn=MeanAbsolutePercentageError(), + optimizer=tf.keras.optimizers.Adam(), + es_patience=100, + save_results=True) + save_results_path = os.path.join(self.path, "model_evaluations_paper") + save_model_path = os.path.join(self.path, "saved_weights") + self.assertTrue(os.path.exists(save_results_path)) + self.assertTrue(os.path.exists(save_model_path)) + + file_path_df = save_results_path+"/model.csv" + df = pd.read_csv(file_path_df) + self.assertEqual(len(df), 1) + for item in [ + "train_loss", "val_loss", "test_loss", + "test_loss_orig", "training_time", + "loss_history", "val_loss_history"]: + self.assertIn(item, df.columns) + + file_path_model = save_model_path+"/model.pickle" + with open(file_path_model, 'rb') as f: + weights_loaded = pickle.load(f) + weights = model.get_weights() + for w1, w2 in zip(weights_loaded, weights): + np.testing.assert_array_equal(w1, w2) + # Clean up + self.fs.remove_object(path_cases) + self.fs.remove_object(os.path.join( + path_mobility, "commuter_mobility_2022.txt")) + self.fs.remove_object(path_mobility) + self.fs.remove_object(save_results_path) + self.fs.remove_object(save_model_path) @patch("os.path.realpath", return_value="/home/") def test_perform_grid_search(self, mock_realpath): @@ -411,6 +456,38 @@ def test_perform_grid_search(self, mock_realpath): self.fs.remove_object(path_mobility) self.fs.remove_object(os.path.join(self.path, "saves")) + def test_scale_data_valid_data(self): + """Test utils.scale_data with valid input and label data.""" + data = { + # 10 samples, 1 day, 5 nodes, 8 groups + "inputs": np.random.rand(10, 1, 8, 5), + "labels": np.random.rand(10, 1, 8, 5) + } + + scaled_inputs, scaled_labels = utils.scale_data(data, True) + + # Check that the scaled data is not equal to the original data + assert not np.allclose( + data["inputs"].transpose(0, 3, 1, 2), scaled_inputs) + assert not np.allclose( + data["labels"].transpose(0, 3, 1, 2), scaled_labels) + + # Check that the scaled data is log-transformed + assert np.allclose(scaled_inputs, np.log1p( + data["inputs"]).transpose(0, 3, 1, 2)) + assert np.allclose(scaled_labels, np.log1p( + data["labels"]).transpose(0, 3, 1, 2)) + + def test_scale_data_invalid_data(self): + """Test utils.scale_data with invalid (non-numeric) data.""" + data = { + "inputs": np.array([["a", "b"], ["c", "d"]]), # Non-numeric data + "labels": np.array([["e", "f"], ["g", "h"]]) + } + + with self.assertRaises(ValueError): + utils.scale_data(data) + if __name__ == '__main__': unittest.main() From 6b3271625c4069de0135d1e61355c35b3756be8a Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 22 Sep 2025 09:02:38 +0200 Subject: [PATCH 10/41] Ading necessary imports --- .../memilio/surrogatemodel_test/test_surrogatemodel_GNN.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index 2a1ba4070f..912953aeb0 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -30,7 +30,7 @@ import pandas as pd import tensorflow as tf import logging - +import spektral # suppress all autograph warnings from tensorflow logging.getLogger("tensorflow").setLevel(logging.ERROR) From b451deda432b317e823ee6b8a05c496321d42635 Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 22 Sep 2025 09:43:18 +0200 Subject: [PATCH 11/41] Fix imports? --- .../memilio/surrogatemodel/GNN/network_architectures.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py index fe1bb033bd..5917a648fd 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py @@ -19,7 +19,7 @@ ############################################################################# import numpy as np import tensorflow as tf - +import spektral import spektral.layers as spektral_layers import spektral.utils.convolution as spektral_convolution From b486d103f618b7c85fa6ba205e1e3cb7ccdf8ddf Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 22 Sep 2025 10:11:48 +0200 Subject: [PATCH 12/41] Add spektral to requirements --- pycode/memilio-surrogatemodel/requirements-dev.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pycode/memilio-surrogatemodel/requirements-dev.txt b/pycode/memilio-surrogatemodel/requirements-dev.txt index d96c76b022..94086d8674 100644 --- a/pycode/memilio-surrogatemodel/requirements-dev.txt +++ b/pycode/memilio-surrogatemodel/requirements-dev.txt @@ -1,3 +1,5 @@ # first support of python 3.11 pyfakefs>=4.6 coverage>=7.0.1 +spektral>=1.0 +tensorflow>=2.0 From ea4cd3199f7e129ca294b4621ce64d4c3319aec4 Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Mon, 22 Sep 2025 13:13:21 +0200 Subject: [PATCH 13/41] undo scale_data-tests --- .../memilio/surrogatemodel_test/test_surrogatemodel_GNN.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index 912953aeb0..c52994eed8 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -456,6 +456,7 @@ def test_perform_grid_search(self, mock_realpath): self.fs.remove_object(path_mobility) self.fs.remove_object(os.path.join(self.path, "saves")) + ''' def test_scale_data_valid_data(self): """Test utils.scale_data with valid input and label data.""" data = { @@ -487,6 +488,7 @@ def test_scale_data_invalid_data(self): with self.assertRaises(ValueError): utils.scale_data(data) + ''' if __name__ == '__main__': From 20274ab2e42511b0f89da87c34ec1068257f43fb Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Wed, 24 Sep 2025 09:56:48 +0200 Subject: [PATCH 14/41] Refactor imports in GNN_utils.py and update test_surrogatemodel_GNN.py to remove commented-out scale_data test --- .../memilio/surrogatemodel/GNN/GNN_utils.py | 12 ++++++++---- .../surrogatemodel_test/test_surrogatemodel_GNN.py | 2 -- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py index 93f0c4851d..a454b28359 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -22,10 +22,10 @@ import os from sklearn.preprocessing import FunctionTransformer -from memilio.epidata import transformMobilityData as tmd -from memilio.epidata import getDataIntoPandasDataFrame as gd -from memilio.simulation.osecir import (ModelGraph, set_edges) -from memilio.epidata import modifyDataframeSeries as mdfs +# memilio.epidata import transformMobilityData as tmd +# from memilio.epidata import getDataIntoPandasDataFrame as gd +# from memilio.simulation.osecir import (ModelGraph, set_edges) +# from memilio.epidata import modifyDataframeSeries as mdfs def remove_confirmed_compartments(dataset_entries, num_groups): @@ -88,6 +88,7 @@ def make_graph(directory, num_regions, countykey_list, models): :models models: List of osecir Model with one model per county. :returns: Graph-ODE model. """ + from memilio.simulation.osecir import (ModelGraph, set_edges) graph = ModelGraph() for i in range(num_regions): graph.add_node(int(countykey_list[i]), models[i]) @@ -102,6 +103,8 @@ def make_graph(directory, num_regions, countykey_list, models): def transform_mobility_directory(): """ Transforms the mobility data by merging Eisenach and Wartburgkreis """ + from memilio.epidata import transformMobilityData as tmd + from memilio.epidata import getDataIntoPandasDataFrame as gd # get mobility data directory arg_dict = gd.cli("commuter_official") @@ -119,6 +122,7 @@ def get_population(): """ Loading the population data for the different counties and the different age groups. """ + from memilio.epidata import modifyDataframeSeries as mdfs df_population = pd.read_json( 'data/Germany/pydata/county_population.json') age_groups = ['0-4', '5-14', '15-34', '35-59', '60-79', '80-130'] diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index c52994eed8..912953aeb0 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -456,7 +456,6 @@ def test_perform_grid_search(self, mock_realpath): self.fs.remove_object(path_mobility) self.fs.remove_object(os.path.join(self.path, "saves")) - ''' def test_scale_data_valid_data(self): """Test utils.scale_data with valid input and label data.""" data = { @@ -488,7 +487,6 @@ def test_scale_data_invalid_data(self): with self.assertRaises(ValueError): utils.scale_data(data) - ''' if __name__ == '__main__': From 35fb82240af57bb1a8bd2c4e8914948900d6ee85 Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Thu, 25 Sep 2025 08:11:05 +0200 Subject: [PATCH 15/41] Add model building and training step to evaluate_and_train.py; update tests --- .../memilio/surrogatemodel/GNN/evaluate_and_train.py | 5 +++++ .../memilio/surrogatemodel/GNN/grid_search.py | 6 +++--- .../surrogatemodel_test/test_surrogatemodel_GNN.py | 8 ++++++-- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index 9ca05be2dc..75149308dc 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -171,6 +171,11 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ train_data, valid_data, test_data = data[:n_train], data[n_train:n_train + n_valid], data[n_train+n_valid:] + # Build the model by calling it on a batch of data + loader = MixedLoader(data) + inputs, _ = loader.__next__() + model(inputs) + # Define Data Loaders loader_tr = MixedLoader( train_data, batch_size=batch_size, epochs=epochs, shuffle=False) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py index 2b2b5850de..0563d98f31 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py @@ -22,8 +22,8 @@ import pandas as pd import tensorflow as tf -from tensorflow.keras.losses import MeanAbsolutePercentageError -from tensorflow.keras.metrics import mean_absolute_percentage_error +from tensorflow.keras.losses import MeanAbsolutePercentageError as MeanAbsolutePercentageError +from tensorflow.keras.metrics import MeanAbsolutePercentageError as mean_absolute_percentage_error from tensorflow.keras.optimizers import Adam from spektral.data import MixedLoader @@ -100,7 +100,7 @@ def perform_grid_search(model_parameters, training_parameters, data): model.compile( optimizer=optimizer, loss=MeanAbsolutePercentageError(), - metrics=[mean_absolute_percentage_error] + metrics=[mean_absolute_percentage_error()] ) results = train_and_evaluate( diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index 912953aeb0..44bccb8e3f 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -280,7 +280,7 @@ def test_train_step(self): def test_evaluate(self): from memilio.surrogatemodel.GNN.evaluate_and_train import ( - evaluate, create_dataset, MixedLoader) + evaluate, create_dataset, MixedLoader, train_step) from memilio.surrogatemodel.GNN.network_architectures import get_model from tensorflow.keras.losses import MeanAbsolutePercentageError @@ -302,6 +302,11 @@ def test_evaluate(self): # Create dataset dataset = create_dataset( path_cases, path_mobility, number_of_nodes=num_nodes) + # Build the model by calling it on a batch of data + loader = MixedLoader(dataset, batch_size=2, epochs=1) + inputs, _ = loader.__next__() + model(inputs) + # Redefine the loader loader = MixedLoader(dataset, batch_size=2, epochs=1) res = evaluate(loader, model, loss_fn) # Check if the result is a tuple of (loss, accuracy) @@ -445,7 +450,6 @@ def test_perform_grid_search(self, mock_realpath): # Check if the results file has the expected number of rows df_results = pd.read_csv(results_file) - print(df_results.columns) self.assertEqual(len(df_results), len(model_parameters)) self.assertEqual(len(df_results.columns), 10) From 9f37f74497d1b892f7f53e56c236c24f4300951d Mon Sep 17 00:00:00 2001 From: Manuel Heger Date: Thu, 25 Sep 2025 09:43:42 +0200 Subject: [PATCH 16/41] Update requirements --- pycode/memilio-surrogatemodel/requirements-dev.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pycode/memilio-surrogatemodel/requirements-dev.txt b/pycode/memilio-surrogatemodel/requirements-dev.txt index 94086d8674..8594db7ac1 100644 --- a/pycode/memilio-surrogatemodel/requirements-dev.txt +++ b/pycode/memilio-surrogatemodel/requirements-dev.txt @@ -1,5 +1,5 @@ # first support of python 3.11 pyfakefs>=4.6 coverage>=7.0.1 -spektral>=1.0 -tensorflow>=2.0 +spektral>=1.2 +tensorflow>=2.12.0 From 68fa1c0449dc2c4040404bd46080af04459e43a5 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 29 Oct 2025 11:01:58 +0100 Subject: [PATCH 17/41] formating and fix tests --- .../GNN/network_architectures.py | 35 +++++++++++++++---- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py index 5917a648fd..d2b662b405 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py @@ -26,7 +26,8 @@ # Function to generate a model class dynamically -def generate_model_class(name, layer_types, num_repeat, num_output, transform=None): +def generate_model_class( + name, layer_types, num_repeat, num_output, transform=None): '''Generates a custom Keras model class with specified layers. :param name: Name of the generated class. @@ -53,10 +54,25 @@ def __init__(self): self.output_layer = tf.keras.layers.Dense( num_output, activation="relu") - def call(self, inputs): + def call(self, inputs, mask=None): # Input consists of a tuple (x, a) where x is the node features and a is the adjacency matrix x, a = inputs + # Extract potential node mask coming from the data loader and ensure broadcastable shape + node_mask = None + if mask is not None: + # When inputs is a list [x, a], Keras/Spektral typically forwards a list of masks [node_mask, None] + node_mask = mask[0] if isinstance(mask, (tuple, list)) else mask + if node_mask is not None: + # Expected by Spektral: shape [B, N, 1] to broadcast with [B, N, C] + if tf.rank(node_mask) == 2: + node_mask = tf.expand_dims(node_mask, axis=-1) + # If no node_mask was provided by the pipeline, create a broadcastable ones-mask + if node_mask is None: + # x has shape [B, N, F] + x_shape = tf.shape(x) + node_mask = tf.ones([x_shape[0], x_shape[1], 1], dtype=tf.float32) + # Apply transformation on adjacency matrix if provided if not tf.is_symbolic_tensor(a): if transform is not None: @@ -65,7 +81,7 @@ def call(self, inputs): # Pass through the layers for layer in self.layer_seq: if type(layer).__module__.startswith("spektral.layers"): - x = layer([x, a]) # Pass both `x` and `a` to the layer + # Pass both `x` and `a` to the layer else: x = layer(x) # Pass only `x` to the layer @@ -90,7 +106,8 @@ def get_model(layer_type, num_layers, num_channels, activation, num_output=1): :param num_output: Number of output units in the final layer. :return: A Keras model instance with the specified architecture. """ - if layer_type not in ["ARMAConv", "GCSConv", "GATConv", "GCNConv", "APPNPConv"]: + if layer_type not in [ + "ARMAConv", "GCSConv", "GATConv", "GCNConv", "APPNPConv"]: raise ValueError( f"Unsupported layer_type: {layer_type}. Supported types are 'ARMAConv', 'GCSConv', 'GATConv', 'GCNConv', 'APPNPConv'.") if num_layers < 1: @@ -141,11 +158,15 @@ def transform(a): a = spektral_convolution.gcn_filter(a) return tf.convert_to_tensor(a, dtype=tf.float32) # Generate the input for generate_model_class - layer_types = [lambda: layer_name( - num_channels, activation=activation, kernel_initializer=tf.keras.initializers.GlorotUniform())] + layer_types = [ + lambda: + layer_name( + num_channels, activation=activation, + kernel_initializer=tf.keras.initializers.GlorotUniform())] num_repeat = [num_layers] model_class = generate_model_class( - "CustomModel", layer_types, num_repeat, num_output, transform=transform) + "CustomModel", layer_types, num_repeat, num_output, + transform=transform) return model_class() From a47f11e83eba876a6337b0fa797e3a959efc8030 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 29 Oct 2025 11:23:48 +0100 Subject: [PATCH 18/41] . --- .../memilio/surrogatemodel/GNN/network_architectures.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py index d2b662b405..8a8b384261 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py @@ -82,6 +82,7 @@ def call(self, inputs, mask=None): for layer in self.layer_seq: if type(layer).__module__.startswith("spektral.layers"): # Pass both `x` and `a` to the layer + x = layer([x, a], mask=[node_mask, None]) else: x = layer(x) # Pass only `x` to the layer From b6f40bb284134e4f8d47a8564741970da2325d87 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Thu, 30 Oct 2025 15:54:54 +0100 Subject: [PATCH 19/41] . --- .../surrogatemodel/GNN/evaluate_and_train.py | 50 ++++++++++--------- .../memilio/surrogatemodel/GNN/grid_search.py | 41 ++++++++------- .../test_surrogatemodel_GNN.py | 30 ++++++----- 3 files changed, 67 insertions(+), 54 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index 75149308dc..5a859bfc70 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -35,8 +35,9 @@ from memilio.surrogatemodel.utils.helper_functions import (calc_split_index) from spektral.data import MixedLoader -from spektral.layers import (GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv, AGNNConv, - APPNPConv, CrystalConv, GATConv, GINConv, XENetConv, GCNConv, GCSConv) +from spektral.layers import ( + GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv, AGNNConv, APPNPConv, + CrystalConv, GATConv, GINConv, XENetConv, GCNConv, GCSConv) from spektral.utils.convolution import normalized_laplacian, rescale_laplacian @@ -150,7 +151,9 @@ def evaluate(loader, model, loss_fn, retransform=False): return np.average(output[:, :-1], 0, weights=output[:, -1]) -def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_patience, save_results=False, save_name="model"): +def train_and_evaluate( + data, batch_size, epochs, model, loss_fn, optimizer, es_patience, + save_dir="", save_name="model"): '''Train and evaluate the GNN model. :param data: The dataset to use for training and evaluation. :param batch_size: Batch size for training. @@ -159,8 +162,8 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ :param loss_fn: Loss function to use. :param optimizer: Optimizer to use for training. :param es_patience: Patience for early stopping. - :param save_results: Whether to save the results to a file. - :param save_name: Name to use when saving the results. + :param save_dir: Directory to save results into. If not provided, results are not saved. + :param save_name: Name to use when saving the results (without extension). :return: A dictionary containing training and evaluation results if save_results is False.''' n = len(data) # Determine split indices for training, validation, and test sets @@ -279,7 +282,7 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ ################################################################################ # Save results ################################################################################ - if save_results: + if save_dir: # save df df.loc[len(df.index)] = [ # layer_name, number_of_layer, channels, np.mean(train_losses), @@ -296,11 +299,9 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ save_name += '.pickle' # Save best weights as pickle - path = os.path.dirname(os.path.realpath(__file__)) - file_path_w = os.path.join(path, 'saved_weights') + file_path_w = os.path.join(save_dir, 'saved_weights') # Ensure the directory exists - if not os.path.isdir(file_path_w): - os.mkdir(file_path_w) + os.makedirs(file_path_w, exist_ok=True) # Construct the full file path by joining the directory with save_name file_path_w = os.path.join(file_path_w, save_name) @@ -309,13 +310,11 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ with open(file_path_w, 'wb') as f: pickle.dump(best_weights, f) + # Save evaluation CSV + file_path_df = os.path.join(save_dir, 'model_evaluations_paper') + os.makedirs(file_path_df, exist_ok=True) file_path_df = os.path.join( - os.path.dirname( - os.path.realpath(os.path.dirname(os.path.realpath(path)))), - 'model_evaluations_paper') - if not os.path.isdir(file_path_df): - os.mkdir(file_path_df) - file_path_df = file_path_df+"/"+save_name.replace('.pickle', '.csv') + file_path_df, save_name.replace('.pickle', '.csv')) df.to_csv(file_path_df) else: @@ -339,9 +338,14 @@ def train_and_evaluate(data, batch_size, epochs, model, loss_fn, optimizer, es_ optimizer = Adam(learning_rate=0.001) loss_fn = MeanAbsolutePercentageError() + path_memilio = os.path.abspath( + os.path.join( + os.path.dirname(__file__), '../../../../..')) + # Generate the Dataset - path_cases = "/localdata1/hege_mn/memilio/saves/GNN_data_30days_3dampings_classic5.pickle" - path_mobility = '/localdata1/hege_mn/memilio/data/Germany/mobility' + path_cases = os.path.join(path_memilio, "saves", + "GNN_data_30days_3dampings_classic5.pickle") + path_mobility = os.path.join(path_memilio, "data", "Germany", "mobility") data = create_dataset(path_cases, path_mobility) # Define the model architecture @@ -362,13 +366,11 @@ def transform_a(adjacency_matrix): model = model_class() - # early stopping patience - # name for df save_name = 'GNN_30days' # name for model - # for param in parameters: - - train_and_evaluate(data, batch_size, epochs, model, - loss_fn, optimizer, es_patience, save_name) + train_and_evaluate( + data, batch_size, epochs, model, loss_fn, optimizer, es_patience, + save_dir=os.path.join(path_memilio, "saves", "GNN_model_results"), + save_name=save_name) elapsed_hyper = time.perf_counter() - start_hyper print( diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py index 0563d98f31..76600bef21 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py @@ -43,10 +43,12 @@ numbers_layers = [2, 3, 4, 5, 6, 7] numbers_channels = [2, 3, 4, 5, 6, 7] activation_functions = ["elu", "relu", "tanh", "sigmoid"] -model_parameters = [(layer, num_layers, num_channels, activation) for layer in layer_types - for num_layers in numbers_layers - for num_channels in numbers_channels - for activation in activation_functions] +model_parameters = [ + (layer, num_layers, num_channels, activation) + for layer in layer_types + for num_layers in numbers_layers + for num_channels in numbers_channels for activation in + activation_functions] # Fix training parameters batch_size = 32 @@ -59,7 +61,8 @@ optimizer, es_patience, max_epochs] -def perform_grid_search(model_parameters, training_parameters, data): +def perform_grid_search( + model_parameters, training_parameters, data, save_dir: str | None = None): """Perform grid search over the specified model parameters. Trains and evaluates models with different configurations and stores the results in a CSV file. Results are saved in 'grid_search_results.csv' in the 'saves' directory. @@ -72,9 +75,13 @@ def perform_grid_search(model_parameters, training_parameters, data): output_dim = data[0].y.shape[-1] batch_size, loss_function, optimizer, es_patience, max_epochs = training_parameters # Create a DataFrame to store the results - df_results = pd.DataFrame(columns=['model', 'optimizer', 'number_of_hidden_layers', 'number_of_channels', 'activation', + df_results = pd.DataFrame( + columns=['model', 'optimizer', + 'number_of_hidden_layers', + 'number_of_channels', 'activation', 'mean_train_loss', - 'mean_validation_loss', 'training_time', 'train_losses', 'val_losses']) + 'mean_validation_loss', 'training_time', + 'train_losses', 'val_losses']) for param in model_parameters: layer_type, num_layers, num_channels, activation = param @@ -106,7 +113,7 @@ def perform_grid_search(model_parameters, training_parameters, data): results = train_and_evaluate( data, batch_size, epochs=max_epochs, model=model, loss_fn=loss_function, optimizer=optimizer, - es_patience=es_patience, save_name="") + es_patience=es_patience, save_name="", save_dir=save_dir) df_results.loc[len(df_results.index)] = [ layer_type, optimizer.__class__.__name__, num_layers, num_channels, activation, @@ -117,17 +124,15 @@ def perform_grid_search(model_parameters, training_parameters, data): results["val_losses"] ] # Save intermediate results to avoid data loss - path = os.path.dirname(os.path.realpath(__file__)) - - file_path_df = os.path.join(os.path.join( - os.path.dirname( - os.path.realpath(os.path.dirname(os.path.realpath(path)))), 'saves')) - if not os.path.isdir(file_path_df): - os.mkdir(file_path_df) - + # If save_dir is provided, use it. Otherwise store under a local 'saves' folder next to this module + base_dir = save_dir if ( + save_dir and len(save_dir) > 0) else os.path.realpath( + os.path.dirname(__file__)) + saves_dir = os.path.join(base_dir, 'saves') + os.makedirs(saves_dir, exist_ok=True) df_results.to_csv(os.path.join( - file_path_df, 'grid_search_results.csv'), index=False) - print(f"Saved intermediate results to {file_path_df}") + saves_dir, 'grid_search_results.csv'), index=False) + print(f"Saved intermediate results to {saves_dir}") # Clear session to free memory after each iteration tf.keras.backend.clear_session() diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index 44bccb8e3f..bb24ace3e7 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -39,7 +39,8 @@ class TestSurrogatemodelGNN(fake_filesystem_unittest.TestCase): path = "/home/" - def create_dummy_data(self, num_samples, num_nodes, num_node_features, output_dim): + def create_dummy_data( + self, num_samples, num_nodes, num_node_features, output_dim): """ Create dummy data for testing. @@ -130,8 +131,9 @@ def test_generate_model_class(self): ModelClass = generate_model_class( "TestModel", layer_types, num_layers, num_output) model = ModelClass() - self.assertEqual(str( - error.exception), "layer_types and num_repeat must have the same length.") + self.assertEqual( + str(error.exception), + "layer_types and num_repeat must have the same length.") # Test with multiple layer types layer_types = [ @@ -176,8 +178,9 @@ def test_get_model(self): with self.assertRaises(ValueError) as error: model = get_model(layer_type, num_layers, num_channels, activation, num_output) - self.assertEqual(str( - error.exception), "Unsupported layer_type: MonvConv. " + self.assertEqual( + str(error.exception), + "Unsupported layer_type: MonvConv. " "Supported types are 'ARMAConv', 'GCSConv', 'GATConv', 'GCNConv', 'APPNPConv'.") # Test with invalud num_layers layer_type = "GATConv" @@ -209,8 +212,9 @@ def test_get_model(self): with self.assertRaises(ValueError) as error: model = get_model(layer_type, num_layers, num_channels, activation, num_output) - self.assertEqual(str( - error.exception), "activation must be a string representing the activation function.") + self.assertEqual( + str(error.exception), + "activation must be a string representing the activation function.") def test_create_dataset(self): from memilio.surrogatemodel.GNN.evaluate_and_train import create_dataset @@ -286,7 +290,8 @@ def test_evaluate(self): # Create a simple model for testing model = get_model( - layer_type="GCNConv", num_layers=2, num_channels=16, activation="relu", num_output=4) + layer_type="GCNConv", num_layers=2, num_channels=16, + activation="relu", num_output=4) loss_fn = MeanAbsolutePercentageError() # Create dummy data in the fake filesystem for testing @@ -337,7 +342,8 @@ def test_train_and_evaluate(self, mock_realpath): number_of_epochs = 2 # Create a simple model for testing model = get_model( - layer_type="GCNConv", num_layers=2, num_channels=16, activation="relu", num_output=4) + layer_type="GCNConv", num_layers=2, num_channels=16, + activation="relu", num_output=4) # Create dummy data in the fake filesystem for testing num_samples = 20 @@ -375,7 +381,7 @@ def test_train_and_evaluate(self, mock_realpath): loss_fn=MeanAbsolutePercentageError(), optimizer=tf.keras.optimizers.Adam(), es_patience=100, - save_results=True) + save_dir=self.path) save_results_path = os.path.join(self.path, "model_evaluations_paper") save_model_path = os.path.join(self.path, "saved_weights") self.assertTrue(os.path.exists(save_results_path)) @@ -430,8 +436,8 @@ def test_perform_grid_search(self, mock_realpath): num_channels = 2 model_parameters = [ - (layer, n_layer, num_channels, activation) for layer in layers for n_layer in num_layers - ] + (layer, n_layer, num_channels, activation) + for layer in layers for n_layer in num_layers] batch_size = 2 loss_function = MeanAbsolutePercentageError() optimizer = tf.keras.optimizers.Adam() From 5bc242ec99d4aacb8fc6fa104ecd2ccac8d4369d Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Thu, 30 Oct 2025 16:23:20 +0100 Subject: [PATCH 20/41] support for py 3.8 --- .../memilio/surrogatemodel/GNN/grid_search.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py index 76600bef21..f92ceccc23 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py @@ -19,6 +19,7 @@ ############################################################################# import os +from typing import Optional import pandas as pd import tensorflow as tf @@ -62,7 +63,8 @@ def perform_grid_search( - model_parameters, training_parameters, data, save_dir: str | None = None): + model_parameters, training_parameters, data, + save_dir: Optional[str] = None): """Perform grid search over the specified model parameters. Trains and evaluates models with different configurations and stores the results in a CSV file. Results are saved in 'grid_search_results.csv' in the 'saves' directory. From 70a9b11e2fccb03365037a6a4df436cc62036516 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Fri, 31 Oct 2025 11:12:48 +0100 Subject: [PATCH 21/41] . --- .../surrogatemodel/GNN/evaluate_and_train.py | 21 +++--- .../test_surrogatemodel_GNN.py | 64 +++++++------------ 2 files changed, 34 insertions(+), 51 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index 5a859bfc70..d9cce1fc5a 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -317,17 +317,16 @@ def train_and_evaluate( file_path_df, save_name.replace('.pickle', '.csv')) df.to_csv(file_path_df) - else: - return { - "model": save_name, - "mean_train_loss": np.mean(train_losses), - "mean_val_loss": np.mean(val_losses), - "mean_test_loss": np.mean(test_scores), - "mean_test_loss_orig": np.mean(test_scores_r), - "training_time": elapsed/60, - "train_losses": losses_history_all, - "val_losses": val_losses_history_all - } + return { + "model": save_name, + "mean_train_loss": np.mean(train_losses), + "mean_val_loss": np.mean(val_losses), + "mean_test_loss": np.mean(test_scores), + "mean_test_loss_orig": np.mean(test_scores_r), + "training_time": elapsed/60, + "train_losses": losses_history_all, + "val_losses": val_losses_history_all + } if __name__ == "__main__": diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index bb24ace3e7..e1bccad32e 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -19,9 +19,6 @@ ############################################################################# from pyfakefs import fake_filesystem_unittest -import memilio.surrogatemodel.GNN.network_architectures as gnn_arch -import memilio.surrogatemodel.GNN.GNN_utils as utils - from unittest.mock import patch import os import unittest @@ -31,6 +28,15 @@ import tensorflow as tf import logging import spektral + +import memilio.surrogatemodel.GNN.network_architectures as gnn_arch +import memilio.surrogatemodel.GNN.GNN_utils as utils + +from memilio.surrogatemodel.GNN.evaluate_and_train import ( + create_dataset, train_and_evaluate, evaluate, train_step, MixedLoader) +from memilio.surrogatemodel.GNN.grid_search import perform_grid_search +from memilio.surrogatemodel.GNN.network_architectures import generate_model_class, get_model +from tensorflow.keras.losses import MeanAbsolutePercentageError # suppress all autograph warnings from tensorflow logging.getLogger("tensorflow").setLevel(logging.ERROR) @@ -87,7 +93,6 @@ def setUp(self): self.setUpPyfakefs() def test_generate_model_class(self): - from memilio.surrogatemodel.GNN.network_architectures import generate_model_class # Test parameters layer_types = [ @@ -152,7 +157,6 @@ def test_generate_model_class(self): self.assertEqual(model.output_layer.units, num_output) def test_get_model(self): - from memilio.surrogatemodel.GNN.network_architectures import get_model # Test parameters layer_type = "GCNConv" @@ -217,7 +221,6 @@ def test_get_model(self): "activation must be a string representing the activation function.") def test_create_dataset(self): - from memilio.surrogatemodel.GNN.evaluate_and_train import create_dataset # Create dummy data in the fake filesystem for testing num_samples = 10 @@ -250,8 +253,6 @@ def test_create_dataset(self): self.fs.remove_object(path_mobility) def test_train_step(self): - from memilio.surrogatemodel.GNN.evaluate_and_train import train_step - from tensorflow.keras.losses import MeanAbsolutePercentageError # Create a simple model for testing model = tf.keras.Sequential([ @@ -283,10 +284,6 @@ def test_train_step(self): self.assertGreaterEqual(loss.numpy(), 0) def test_evaluate(self): - from memilio.surrogatemodel.GNN.evaluate_and_train import ( - evaluate, create_dataset, MixedLoader, train_step) - from memilio.surrogatemodel.GNN.network_architectures import get_model - from tensorflow.keras.losses import MeanAbsolutePercentageError # Create a simple model for testing model = get_model( @@ -334,10 +331,6 @@ def test_evaluate(self): @patch("os.path.realpath", return_value="/home/") def test_train_and_evaluate(self, mock_realpath): - from memilio.surrogatemodel.GNN.evaluate_and_train import ( - train_and_evaluate, create_dataset, MixedLoader) - from memilio.surrogatemodel.GNN.network_architectures import get_model - from tensorflow.keras.losses import MeanAbsolutePercentageError number_of_epochs = 2 # Create a simple model for testing @@ -410,16 +403,11 @@ def test_train_and_evaluate(self, mock_realpath): self.fs.remove_object(save_results_path) self.fs.remove_object(save_model_path) - @patch("os.path.realpath", return_value="/home/") - def test_perform_grid_search(self, mock_realpath): - from memilio.surrogatemodel.GNN.evaluate_and_train import ( - create_dataset) - from memilio.surrogatemodel.GNN.grid_search import perform_grid_search - from tensorflow.keras.losses import MeanAbsolutePercentageError + def test_perform_grid_search(self): # Create dummy data in the fake filesystem for testing - num_samples = 20 - num_nodes = 5 + num_samples = 10 + num_nodes = 4 num_node_features = 3 output_dim = 4 data = self.create_dummy_data( @@ -430,14 +418,17 @@ def test_perform_grid_search(self, mock_realpath): # Create dataset dataset = create_dataset( path_cases, path_mobility, number_of_nodes=num_nodes) - layers = ["GCNConv", "GATConv"] - num_layers = [1, 2] - activation = "relu" - num_channels = 2 - model_parameters = [ - (layer, n_layer, num_channels, activation) - for layer in layers for n_layer in num_layers] + # Define model parameters for grid search + layers = ["GCNConv"] + num_layers = [1] + num_channels = [8] + activations = ["relu"] + model_parameters = [(layer, n_layer, channel, activation) + for layer in layers for n_layer in num_layers + for channel in num_channels + for activation in activations + for layer in layers for n_layer in num_layers] batch_size = 2 loss_function = MeanAbsolutePercentageError() optimizer = tf.keras.optimizers.Adam() @@ -445,9 +436,9 @@ def test_perform_grid_search(self, mock_realpath): max_epochs = 2 training_parameters = [batch_size, loss_function, optimizer, es_patience, max_epochs] - # Perform grid search + # Perform grid search with explicit save_dir to avoid os.path.realpath issues perform_grid_search(model_parameters, training_parameters, - dataset) + dataset, save_dir=self.path) # Check if the results file is created results_file = os.path.join( @@ -459,13 +450,6 @@ def test_perform_grid_search(self, mock_realpath): self.assertEqual(len(df_results), len(model_parameters)) self.assertEqual(len(df_results.columns), 10) - # Clean up - self.fs.remove_object(path_cases) - self.fs.remove_object(os.path.join( - path_mobility, "commuter_mobility_2022.txt")) - self.fs.remove_object(path_mobility) - self.fs.remove_object(os.path.join(self.path, "saves")) - def test_scale_data_valid_data(self): """Test utils.scale_data with valid input and label data.""" data = { From 4d42a2a2272fb6db031bb99965b821de79a806dd Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Tue, 4 Nov 2025 12:54:05 +0100 Subject: [PATCH 22/41] [ci skip] start rework data gemeration gnn --- .../surrogatemodel/GNN/data_generation.py | 352 +++++++----------- 1 file changed, 144 insertions(+), 208 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index 1b3269b0bc..a635ee0a16 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -29,12 +29,13 @@ from progress.bar import Bar -from memilio.simulation import (AgeGroup, LogLevel, set_log_level, Damping) -from memilio.simulation.osecir import (Index_InfectionState, interpolate_simulation_result, ParameterStudy, - InfectionState, Model, interpolate_simulation_result) +from memilio.simulation import (AgeGroup, set_log_level, Damping) +from memilio.simulation.osecir import ( + Index_InfectionState, interpolate_simulation_result, ParameterStudy, + InfectionState, Model, interpolate_simulation_result) -from memilio.surrogatemodel.GNN.GNN_utils import (transform_mobility_directory, - make_graph, scale_data, getBaselineMatrix, remove_confirmed_compartments) +from memilio.surrogatemodel.GNN.GNN_utils import ( + scale_data, remove_confirmed_compartments) import memilio.surrogatemodel.utils.dampings as dampings from enum import Enum @@ -48,16 +49,20 @@ class Location(Enum): Other = 3 -# Define the start and the latest end date for the simulation -start_date = mio.Date(2019, 1, 1) -end_date = mio.Date(2021, 12, 31) +# Enumerate the different initialization strategies for compartment populations +class InitializationStrategy(Enum): + """Strategy for initializing compartment populations in the simulation.""" + GROUND_TRUTH_NOISY = "ground_truth_noisy" # Uses ground truth data with noise + BOUNDS_RANDOM = "bounds_random" # Uses random values between bounds def set_covid_parameters(model, num_groups=6): """Setting COVID-parameters for the different age groups. - :param model: memilio model, whose parameters should be clarified - :param num_groups: Number of age groups + Parameters are based on Kühn et al. https://doi.org/10.1016/j.mbs.2021.108648. + + :param model: memilio model, whose parameters should be clarified. + :param num_groups: Number of age groups. """ # age specific parameters @@ -73,11 +78,12 @@ def set_covid_parameters(model, num_groups=6): TimeInfectedSevere = [5, 5, 5.925, 7.55, 8.5, 11] TimeInfectedCritical = [6.95, 6.95, 6.86, 17.36, 17.1, 11.6] - for i, rho, muCR, muHI, muUH, muDU, tc, ti, th, tu in zip(range(num_groups), - TransmissionProbabilityOnContact, RecoveredPerInfectedNoSymptoms, - SeverePerInfectedSymptoms, CriticalPerSevere, DeathsPerCritical, - TimeInfectedNoSymptoms, TimeInfectedSymptoms, - TimeInfectedSevere, TimeInfectedCritical): + for i, rho, muCR, muHI, muUH, muDU, tc, ti, th, tu in zip( + range(num_groups), + TransmissionProbabilityOnContact, RecoveredPerInfectedNoSymptoms, + SeverePerInfectedSymptoms, CriticalPerSevere, DeathsPerCritical, + TimeInfectedNoSymptoms, TimeInfectedSymptoms, TimeInfectedSevere, + TimeInfectedCritical): # Compartment transition duration model.parameters.TimeExposed[AgeGroup(i)] = 3.335 model.parameters.TimeInfectedNoSymptoms[AgeGroup(i)] = tc @@ -100,11 +106,11 @@ def set_covid_parameters(model, num_groups=6): def set_contact_matrices(model, data_dir, num_groups=6): - """Setting the contact matrices for a model + """Setting the contact matrices for a model. :param model: memilio ODE-model, whose contact matrices should be modified. - :param data_dir: directory, where the contact data is stored (should contain folder "contacts") - :param num_groups: Number of age groups considered + :param data_dir: directory, where the contact data is stored (should contain folder "contacts"). + :param num_groups: Number of age groups considered. """ @@ -123,12 +129,17 @@ def set_contact_matrices(model, data_dir, num_groups=6): model.parameters.ContactPatterns.cont_freq_mat = contact_matrices -def get_graph(num_groups, data_dir, mobility_directory): +def get_graph( + num_groups, data_dir, mobility_directory, start_date=mio.Date( + 2020, 6, 1), + end_date=mio.Date(2020, 8, 31)): """ Generate the associated graph given the mobility data. - :param num_groups: Number of age groups - :param data_dir: Directory, where the contact data is stored (should contain folder "contacts") - :param mobility_directory: Directory containing the mobility data + :param num_groups: Number of age groups. + :param data_dir: Directory, where the contact data is stored (should contain folder "contacts"). + :param mobility_directory: Directory containing the mobility data. + :param start_date: Date when the simulation starts. + :param end_date: Date when the simulation ends. """ # Generating and Initializing the model model = Model(num_groups) @@ -150,148 +161,44 @@ def get_graph(num_groups, data_dir, mobility_directory): path_population_data = os.path.join(pydata_dir, "county_current_population.json") - # Setting node information based on model parameters + is_node_for_county = True + mio.osecir.set_nodes( - model.parameters, - mio.Date(start_date.year, - start_date.month, start_date.day), - mio.Date(end_date.year, - end_date.month, end_date.day), pydata_dir, - path_population_data, True, graph, scaling_factor_infected, - scaling_factor_icu, tnt_capacity_factor, 0, False) - - # Setting edge information based on the mobility data + model.parameters, mio.Date( + start_date.year, start_date.month, start_date.day), + mio.Date(end_date.year, end_date.month, end_date.day), + pydata_dir, path_population_data, is_node_for_county, graph, + scaling_factor_infected, scaling_factor_icu, tnt_capacity_factor, 0, False) + mio.osecir.set_edges( mobility_directory, graph, len(Location)) return graph -def run_secir_groups_simulation1(days, damping_days, damping_factors, graph, num_groups=6, start_date=mio.Date(2020, 6, 1)): - """ Uses an ODE SECIR model allowing for asymptomatic infection with 6 - different age groups. The model is not stratified by region. - Virus-specific parameters are fixed and initial number of person - in the particular infection states are chosen randomly from defined ranges. - - :param days: Number of days simulated within a single run. - :param damping_days: Days, where a damping is applied - :param damping_factors: damping factors associated to the damping days - :param graph: Graph initialized for the start_date with the population data which - is sampled during the run. - :param num_groups: Number of age groups considered in the simulation - :param start_date: Date, when the simulation starts - :returns: List containing the populations in each compartment used to initialize - the run. - """ - if len(damping_days) != len(damping_factors): - raise ValueError("Length of damping_days and damping_factors differ!") - - min_date = mio.Date(2020, 6, 1) - min_date_num = min_date.day_in_year - start_date_num = start_date.day_in_year - min_date_num - - # Load the ground truth data - pydata_dir = os.path.join(data_dir, "Germany", "pydata") - ground_truth_dir = os.path.join( - pydata_dir, "ground_truth_all_nodes.pickle") - with open(ground_truth_dir, 'rb') as f: - ground_truth_all_nodes = pickle.load(f) - - # Initialize model for each node, using the population data and sampling the number of - # individuals in the different compartments - for node_indx in range(graph.num_nodes): - model = graph.get_node(node_indx).property - data = ground_truth_all_nodes[node_indx][start_date_num] - - # Iterate over the different age groups - for i in range(num_groups): - age_group = AgeGroup(i) - pop_age_group = model.populations.get_group_total_AgeGroup( - age_group) - - # Generating valid, noisy configuration of the compartments - valid_configuration = False - while valid_configuration is False: - comp_data = data[:, i] - ratios = np.random.uniform(0.8, 1.2, size=8) - init_data = comp_data * ratios - if np.sum(init_data[1:]) < pop_age_group: - valid_configuration = True - - # Set the populations for the different compartments - model.populations[age_group, Index_InfectionState( - InfectionState.Exposed)] = init_data[1] - model.populations[age_group, Index_InfectionState( - InfectionState.InfectedNoSymptoms)] = init_data[2] - model.populations[age_group, Index_InfectionState( - InfectionState.InfectedSymptoms)] = init_data[3] - model.populations[age_group, Index_InfectionState( - InfectionState.InfectedSevere)] = init_data[4] - model.populations[age_group, Index_InfectionState( - InfectionState.InfectedCritical)] = init_data[5] - model.populations[age_group, Index_InfectionState( - InfectionState.Recovered)] = init_data[6] - model.populations[age_group, Index_InfectionState( - InfectionState.Dead)] = init_data[7] - model.populations.set_difference_from_group_total_AgeGroup(( - age_group, InfectionState.Susceptible), pop_age_group) - - # Introduce the damping information, in general dampings can be local, but till now the code just allows global dampings - damped_matrices = [] - damping_coefficients = [] - - for i in np.arange(len(damping_days)): - day = damping_days[i] - factor = damping_factors[i] - - damping = np.ones((num_groups, num_groups) - ) * np.float16(factor) - model.parameters.ContactPatterns.cont_freq_mat.add_damping(Damping( - coeffs=(damping), t=day, level=0, type=0)) - damped_matrices.append(model.parameters.ContactPatterns.cont_freq_mat.get_matrix_at( - day+1)) - damping_coefficients.append(damping) - - # Apply mathematical constraints to parameters - model.apply_constraints() - - # set model to graph - graph.get_node(node_indx).property.populations = model.populations - - # Start simulation - study = ParameterStudy(graph, 0, days, dt=0.5, num_runs=1) - start_time = time.perf_counter() - study.run() - runtime = time.perf_counter() - start_time - - graph_run = study.run()[0] - results = interpolate_simulation_result(graph_run) +def run_secir_groups_simulation( + days, damping_days, damping_factors, graph, num_groups=6, + start_date=mio.Date(2020, 6, 1), end_date=mio.Date(2020, 8, 31), + initialization_strategy=InitializationStrategy.BOUNDS_RANDOM): + """Run an ODE SECIR simulation with multiple age groups and regions. - for result_indx in range(len(results)): - results[result_indx] = remove_confirmed_compartments( - np.asarray(results[result_indx]), num_groups) - - dataset_entry = copy.deepcopy(results) - - return dataset_entry, damped_matrices, damping_coefficients, runtime - - -def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_groups=6, start_date=mio.Date(2020, 6, 1)): - """ Uses an ODE SECIR model allowing for asymptomatic infection with 6 - different age groups. The model is not stratified by region. - Virus-specific parameters are fixed and initial number of person - in the particular infection states are chosen randomly from defined ranges. + Uses an ODE SECIR model allowing for asymptomatic infection with 6 + different age groups. The model is stratified by region using a graph structure. + Virus-specific parameters are fixed and initial number of persons in the + particular infection states are chosen using the specified strategy. :param days: Number of days simulated within a single run. - :param damping_days: Days, where a damping is applied - :param damping_factors: damping factors associated to the damping days - :param graph: Graph initialized for the start_date with the population data which - is sampled during the run. - :param num_groups: Number of age groups considered in the simulation - :param start_date: Date, when the simulation starts - :returns: List containing the populations in each compartment used to initialize - the run. - """ + :param damping_days: List of days where dampings are applied. + :param damping_factors: List of damping factors associated to the damping days. + :param graph: Graph initialized for the start_date with the population data. + :param num_groups: Number of age groups considered in the simulation. + :param start_date: Date when the simulation starts. + :param end_date: Date when the simulation ends. + :param initialization_strategy: Strategy for initializing compartment populations. + - GROUND_TRUTH_NOISY: Uses ground truth data with random noise (ratios 0.8-1.2) + - BOUNDS_RANDOM: Uses random values between predefined bounds + :returns: Tuple containing (simulation_results, damped_matrices, damping_coefficients, runtime) + """ if len(damping_days) != len(damping_factors): raise ValueError("Length of damping_days and damping_factors differ!") @@ -299,24 +206,30 @@ def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_ min_date_num = min_date.day_in_year start_date_num = start_date.day_in_year - min_date_num - # Load the ground truth data + # Load the appropriate ground truth data based on initialization strategy pydata_dir = os.path.join(data_dir, "Germany", "pydata") - upper_bound_dir = os.path.join( - pydata_dir, "ground_truth_upper_bound.pickle") - lower_bound_dir = os.path.join( - pydata_dir, "ground_truth_lower_bound.pickle") - with open(upper_bound_dir, 'rb') as f: - ground_truth_upper_bound = pickle.load(f) - - with open(lower_bound_dir, 'rb') as f: - ground_truth_lower_bound = pickle.load(f) - - # Initialize model for each node, using the population data and sampling the number of - # individuals in the different compartments + + if initialization_strategy == InitializationStrategy.GROUND_TRUTH_NOISY: + ground_truth_dir = os.path.join( + pydata_dir, "ground_truth_all_nodes.pickle") + with open(ground_truth_dir, 'rb') as f: + ground_truth_data = pickle.load(f) + elif initialization_strategy == InitializationStrategy.BOUNDS_RANDOM: + upper_bound_dir = os.path.join( + pydata_dir, "ground_truth_upper_bound.pickle") + lower_bound_dir = os.path.join( + pydata_dir, "ground_truth_lower_bound.pickle") + with open(upper_bound_dir, 'rb') as f: + ground_truth_upper_bound = pickle.load(f) + with open(lower_bound_dir, 'rb') as f: + ground_truth_lower_bound = pickle.load(f) + else: + raise ValueError( + f"Unknown initialization strategy: {initialization_strategy}") + + # Initialize model for each node for node_indx in range(graph.num_nodes): model = graph.get_node(node_indx).property - max_data = ground_truth_upper_bound[node_indx] - min_data = ground_truth_lower_bound[node_indx] # Iterate over the different age groups for i in range(num_groups): @@ -324,14 +237,24 @@ def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_ pop_age_group = model.populations.get_group_total_AgeGroup( age_group) - # Generating valid, noisy configuration of the compartments + # Generate initial compartment populations based on strategy valid_configuration = False while valid_configuration is False: - init_data = np.asarray([0 for _ in range(8)]) - max_val = max_data[:, i] - min_val = min_data[:, i] - for j in range(1, 8): - init_data[j] = random.uniform(min_val[j], max_val[j]) + if initialization_strategy == InitializationStrategy.GROUND_TRUTH_NOISY: + # Use ground truth data with noise + comp_data = ground_truth_data[node_indx][start_date_num][ + :, i] + ratios = np.random.uniform(0.8, 1.2, size=8) + init_data = comp_data * ratios + elif initialization_strategy == InitializationStrategy.BOUNDS_RANDOM: + # Use random values between bounds + init_data = np.asarray([0 for _ in range(8)]) + max_data = ground_truth_upper_bound[node_indx][:, i] + min_data = ground_truth_lower_bound[node_indx][:, i] + for j in range(1, 8): + init_data[j] = random.uniform(min_data[j], max_data[j]) + + # Check if configuration is valid (infected population < total population) if np.sum(init_data[1:]) < pop_age_group: valid_configuration = True @@ -353,7 +276,7 @@ def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_ model.populations.set_difference_from_group_total_AgeGroup(( age_group, InfectionState.Susceptible), pop_age_group) - # Introduce the damping information, in general dampings can be local, but till now the code just allows global dampings + # Apply damping information (currently only global dampings supported) damped_matrices = [] damping_coefficients = [] @@ -361,26 +284,27 @@ def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_ day = damping_days[i] factor = damping_factors[i] - damping = np.ones((num_groups, num_groups) - ) * np.float16(factor) + damping = np.ones((num_groups, num_groups)) * np.float16(factor) model.parameters.ContactPatterns.cont_freq_mat.add_damping(Damping( coeffs=(damping), t=day, level=0, type=0)) - damped_matrices.append(model.parameters.ContactPatterns.cont_freq_mat.get_matrix_at( - day+1)) + damped_matrices.append( + model.parameters.ContactPatterns.cont_freq_mat.get_matrix_at( + day + 1)) damping_coefficients.append(damping) # Apply mathematical constraints to parameters model.apply_constraints() - # set model to graph + # Set model to graph graph.get_node(node_indx).property.populations = model.populations - # Start simulation + # Run simulation study = ParameterStudy(graph, 0, days, dt=0.5, num_runs=1) start_time = time.perf_counter() study.run() runtime = time.perf_counter() - start_time + # Process results graph_run = study.run()[0] results = interpolate_simulation_result(graph_run) @@ -394,21 +318,26 @@ def run_secir_groups_simulation(days, damping_days, damping_factors, graph, num_ def generate_data( - num_runs, data_dir, path, input_width, label_width, save_data=True, - transform=True, damping_method="classic", max_number_damping=3): - """ Generate dataset by calling run_secir_simulation (num_runs)-often + num_runs, data_dir, path, input_width, label_width, start_date, + end_date, save_data=True, transform=True, damping_method="classic", + max_number_damping=3, + initialization_strategy=InitializationStrategy.BOUNDS_RANDOM): + """Generate dataset by calling run_secir_groups_simulation multiple times. - :param num_runs: Number of times, the function run_secir_simulation is called. + :param num_runs: Number of simulation runs to generate. :param data_dir: Directory with all data needed to initialize the models. - :param path: Path, where the datasets are stored. - :param input_width: number of time steps used for model input. - :param label_width: number of time steps (days) used as model output/label. - :param save_data: Option to deactivate the save of the dataset. Per default True. - :param transform: Option to deactivate the transformation of the data. Per default True. - :param damping_method: String specifying the damping method, that should be used. Possible values "classic", "active", "random". - :param max_number_damping: Maximal number of possible dampings. - :returns: Dictionary containing inputs, labels, contact_matrix, damping_days and damping_factors - """ + :param path: Path where the datasets are stored. + :param input_width: Number of time steps used for model input. + :param label_width: Number of time steps (days) used as model output/label. + :param start_date: Date when the simulation starts. + :param end_date: Date when the simulation ends. + :param save_data: Option to deactivate the save of the dataset. Default True. + :param transform: Option to deactivate the transformation of the data. Default True. + :param damping_method: String specifying the damping method. Values: "classic", "active", "random". + :param max_number_damping: Maximal number of possible dampings. + :param initialization_strategy: Strategy for initializing compartment populations. + :returns: Dictionary containing inputs, labels, contact_matrix, damping_days and damping_factors. + """ set_log_level(mio.LogLevel.Error) days = label_width + input_width - 1 @@ -446,14 +375,15 @@ def generate_data( # Generate random damping days and damping factors if max_number_damping > 0: damping_days, damping_factors = dampings.generate_dampings( - days, max_number_damping, method=damping_method, min_distance=2, - min_damping_day=2) + days, max_number_damping, method=damping_method, + min_distance=2, min_damping_day=2) else: damping_days = [] damping_factors = [] # Run simulation data_run, damped_matrices, damping_coefficients, t_run = run_secir_groups_simulation( - days, damping_days, damping_factors, graph, num_groups, start_date) + days, damping_days, damping_factors, graph, num_groups, start_date, + initialization_strategy) times.append(t_run) @@ -507,20 +437,26 @@ def generate_data( path = os.getcwd() path_output = os.path.join(os.getcwd(), 'saves') - # data_dir = os.path.join(os.getcwd(), 'data') - data_dir = "/localdata1/hege_mn/memilio/data" + data_dir = os.path.join(os.getcwd(), 'data') input_width = 5 number_of_dampings = 0 num_runs = 1 label_width_list = [30] + # Define the start and the latest end date for the simulation + start_date = mio.Date(2019, 1, 1) + end_date = mio.Date(2021, 12, 31) + random.seed(10) + + # Other option is InitializationStrategy.GROUND_TRUTH_NOISY + init_strategy = InitializationStrategy.BOUNDS_RANDOM + + # init_strategy = for label_width in label_width_list: - generate_data(num_runs=num_runs, - data_dir=data_dir, - path=path_output, - input_width=input_width, - label_width=label_width, - save_data=True, - damping_method="active", - max_number_damping=number_of_dampings) + generate_data( + num_runs=num_runs, data_dir=data_dir, path=path_output, + input_width=input_width, label_width=label_width, + start_date=start_date, end_date=end_date, save_data=True, + damping_method="active", max_number_damping=number_of_dampings, + initialization_strategy=init_strategy) From 8c77ef25a861c37760b8f68dc5c70d5d0dca8a9b Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Tue, 4 Nov 2025 13:06:48 +0100 Subject: [PATCH 23/41] [ci skip] . --- .../surrogatemodel/GNN/data_generation.py | 109 +++++------------- 1 file changed, 27 insertions(+), 82 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index a635ee0a16..7cd5217ac8 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -49,13 +49,6 @@ class Location(Enum): Other = 3 -# Enumerate the different initialization strategies for compartment populations -class InitializationStrategy(Enum): - """Strategy for initializing compartment populations in the simulation.""" - GROUND_TRUTH_NOISY = "ground_truth_noisy" # Uses ground truth data with noise - BOUNDS_RANDOM = "bounds_random" # Uses random values between bounds - - def set_covid_parameters(model, num_groups=6): """Setting COVID-parameters for the different age groups. @@ -131,8 +124,7 @@ def set_contact_matrices(model, data_dir, num_groups=6): def get_graph( num_groups, data_dir, mobility_directory, start_date=mio.Date( - 2020, 6, 1), - end_date=mio.Date(2020, 8, 31)): + 2020, 6, 1), end_date=mio.Date(2020, 8, 31)): """ Generate the associated graph given the mobility data. :param num_groups: Number of age groups. @@ -177,15 +169,13 @@ def get_graph( def run_secir_groups_simulation( - days, damping_days, damping_factors, graph, num_groups=6, - start_date=mio.Date(2020, 6, 1), end_date=mio.Date(2020, 8, 31), - initialization_strategy=InitializationStrategy.BOUNDS_RANDOM): + days, damping_days, damping_factors, graph, num_groups=6): """Run an ODE SECIR simulation with multiple age groups and regions. Uses an ODE SECIR model allowing for asymptomatic infection with 6 different age groups. The model is stratified by region using a graph structure. Virus-specific parameters are fixed and initial number of persons in the - particular infection states are chosen using the specified strategy. + particular infection states are chosen randomly between predefined bounds. :param days: Number of days simulated within a single run. :param damping_days: List of days where dampings are applied. @@ -193,39 +183,21 @@ def run_secir_groups_simulation( :param graph: Graph initialized for the start_date with the population data. :param num_groups: Number of age groups considered in the simulation. :param start_date: Date when the simulation starts. - :param end_date: Date when the simulation ends. - :param initialization_strategy: Strategy for initializing compartment populations. - - GROUND_TRUTH_NOISY: Uses ground truth data with random noise (ratios 0.8-1.2) - - BOUNDS_RANDOM: Uses random values between predefined bounds :returns: Tuple containing (simulation_results, damped_matrices, damping_coefficients, runtime) """ if len(damping_days) != len(damping_factors): raise ValueError("Length of damping_days and damping_factors differ!") - min_date = mio.Date(2020, 6, 1) - min_date_num = min_date.day_in_year - start_date_num = start_date.day_in_year - min_date_num - - # Load the appropriate ground truth data based on initialization strategy + # Load ground truth bounds for initialization pydata_dir = os.path.join(data_dir, "Germany", "pydata") - - if initialization_strategy == InitializationStrategy.GROUND_TRUTH_NOISY: - ground_truth_dir = os.path.join( - pydata_dir, "ground_truth_all_nodes.pickle") - with open(ground_truth_dir, 'rb') as f: - ground_truth_data = pickle.load(f) - elif initialization_strategy == InitializationStrategy.BOUNDS_RANDOM: - upper_bound_dir = os.path.join( - pydata_dir, "ground_truth_upper_bound.pickle") - lower_bound_dir = os.path.join( - pydata_dir, "ground_truth_lower_bound.pickle") - with open(upper_bound_dir, 'rb') as f: - ground_truth_upper_bound = pickle.load(f) - with open(lower_bound_dir, 'rb') as f: - ground_truth_lower_bound = pickle.load(f) - else: - raise ValueError( - f"Unknown initialization strategy: {initialization_strategy}") + upper_bound_dir = os.path.join( + pydata_dir, "ground_truth_upper_bound.pickle") + lower_bound_dir = os.path.join( + pydata_dir, "ground_truth_lower_bound.pickle") + with open(upper_bound_dir, 'rb') as f: + ground_truth_upper_bound = pickle.load(f) + with open(lower_bound_dir, 'rb') as f: + ground_truth_lower_bound = pickle.load(f) # Initialize model for each node for node_indx in range(graph.num_nodes): @@ -237,22 +209,14 @@ def run_secir_groups_simulation( pop_age_group = model.populations.get_group_total_AgeGroup( age_group) - # Generate initial compartment populations based on strategy + # Generate initial compartment populations with random values between bounds valid_configuration = False while valid_configuration is False: - if initialization_strategy == InitializationStrategy.GROUND_TRUTH_NOISY: - # Use ground truth data with noise - comp_data = ground_truth_data[node_indx][start_date_num][ - :, i] - ratios = np.random.uniform(0.8, 1.2, size=8) - init_data = comp_data * ratios - elif initialization_strategy == InitializationStrategy.BOUNDS_RANDOM: - # Use random values between bounds - init_data = np.asarray([0 for _ in range(8)]) - max_data = ground_truth_upper_bound[node_indx][:, i] - min_data = ground_truth_lower_bound[node_indx][:, i] - for j in range(1, 8): - init_data[j] = random.uniform(min_data[j], max_data[j]) + init_data = np.asarray([0 for _ in range(8)]) + max_data = ground_truth_upper_bound[node_indx][:, i] + min_data = ground_truth_lower_bound[node_indx][:, i] + for j in range(1, 8): + init_data[j] = random.uniform(min_data[j], max_data[j]) # Check if configuration is valid (infected population < total population) if np.sum(init_data[1:]) < pop_age_group: @@ -320,8 +284,7 @@ def run_secir_groups_simulation( def generate_data( num_runs, data_dir, path, input_width, label_width, start_date, end_date, save_data=True, transform=True, damping_method="classic", - max_number_damping=3, - initialization_strategy=InitializationStrategy.BOUNDS_RANDOM): + max_number_damping=3, mobility_file="commuter_mobility_2022.txt"): """Generate dataset by calling run_secir_groups_simulation multiple times. :param num_runs: Number of simulation runs to generate. @@ -335,7 +298,7 @@ def generate_data( :param transform: Option to deactivate the transformation of the data. Default True. :param damping_method: String specifying the damping method. Values: "classic", "active", "random". :param max_number_damping: Maximal number of possible dampings. - :param initialization_strategy: Strategy for initializing compartment populations. + :param mobility_file: Filename of the mobility data file. :returns: Dictionary containing inputs, labels, contact_matrix, damping_days and damping_factors. """ set_log_level(mio.LogLevel.Error) @@ -352,18 +315,8 @@ def generate_data( # Setting basic parameter num_groups = 6 - mobility_dir = data_dir + "/Germany/mobility/commuter_mobility_2022.txt" - graph = get_graph(num_groups, data_dir, mobility_dir) - # Define possible start dates for the simulation - # start_dates = [ - # mio.Date(2020, 6, 1), - # mio.Date(2020, 7, 1), - # mio.Date(2020, 8, 1), - # mio.Date(2020, 9, 1), - # mio.Date(2020, 10, 1), - # mio.Date(2020, 11, 1), - # mio.Date(2020, 12, 1) - # ] + mobility_dir = data_dir + "/Germany/mobility/" + mobility_file + graph = get_graph(num_groups, data_dir, mobility_dir, start_date, end_date) # show progess in terminal for longer runs # Due to the random structure, there is currently no need to shuffle the data @@ -371,7 +324,6 @@ def generate_data( times = [] for i in range(0, num_runs): - # start_date = start_dates[i % len(start_dates)] # Generate random damping days and damping factors if max_number_damping > 0: damping_days, damping_factors = dampings.generate_dampings( @@ -382,8 +334,7 @@ def generate_data( damping_factors = [] # Run simulation data_run, damped_matrices, damping_coefficients, t_run = run_secir_groups_simulation( - days, damping_days, damping_factors, graph, num_groups, start_date, - initialization_strategy) + days, damping_days, damping_factors, graph, num_groups) times.append(t_run) @@ -442,21 +393,15 @@ def generate_data( number_of_dampings = 0 num_runs = 1 label_width_list = [30] - - # Define the start and the latest end date for the simulation - start_date = mio.Date(2019, 1, 1) - end_date = mio.Date(2021, 12, 31) - random.seed(10) - # Other option is InitializationStrategy.GROUND_TRUTH_NOISY - init_strategy = InitializationStrategy.BOUNDS_RANDOM + # Define the start and the latest end date for the simulation + start_date = mio.Date(2020, 10, 1) + end_date = mio.Date(2021, 10, 31) - # init_strategy = for label_width in label_width_list: generate_data( num_runs=num_runs, data_dir=data_dir, path=path_output, input_width=input_width, label_width=label_width, start_date=start_date, end_date=end_date, save_data=True, - damping_method="active", max_number_damping=number_of_dampings, - initialization_strategy=init_strategy) + damping_method="active", max_number_damping=number_of_dampings) From 4d72f231c8a188b9a181623a7d29be54348b8b4d Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 11:01:15 +0100 Subject: [PATCH 24/41] [ci skip] complete rework data generation --- .../surrogatemodel/GNN/data_generation.py | 745 +++++++++++------- 1 file changed, 480 insertions(+), 265 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index 7cd5217ac8..c91a58849d 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -17,294 +17,441 @@ # See the License for the specific language governing permissions and # limitations under the License. ############################################################################# +""" +Data generation module for GNN-based surrogate models. + +This module provides functionality to generate training data for Graph Neural Network +surrogate models by running a ODE-SECIR model based simulations across multiple regions and +time periods with varying damping interventions. +""" import copy import os import pickle import random import time -import memilio.simulation as mio -import memilio.simulation.osecir as osecir -import numpy as np +from typing import Dict, List, Tuple +import numpy as np from progress.bar import Bar -from memilio.simulation import (AgeGroup, set_log_level, Damping) +import memilio.simulation as mio +import memilio.simulation.osecir as osecir +from memilio.simulation import AgeGroup, set_log_level, Damping from memilio.simulation.osecir import ( - Index_InfectionState, interpolate_simulation_result, ParameterStudy, - InfectionState, Model, interpolate_simulation_result) + Index_InfectionState, InfectionState, ParameterStudy, + interpolate_simulation_result) from memilio.surrogatemodel.GNN.GNN_utils import ( - scale_data, remove_confirmed_compartments) + scale_data, + remove_confirmed_compartments +) import memilio.surrogatemodel.utils.dampings as dampings -from enum import Enum - -# Enumerate the different locations class Location(Enum): + """Contact location types for the model.""" Home = 0 School = 1 Work = 2 Other = 3 -def set_covid_parameters(model, num_groups=6): - """Setting COVID-parameters for the different age groups. +# Default number of age groups +DEFAULT_NUM_AGE_GROUPS = 6 - Parameters are based on Kühn et al. https://doi.org/10.1016/j.mbs.2021.108648. +# Contact location names corresponding to provided contact files +CONTACT_LOCATIONS = ["home", "school_pf_eig", "work", "other"] - :param model: memilio model, whose parameters should be clarified. - :param num_groups: Number of age groups. - """ - # age specific parameters - TransmissionProbabilityOnContact = [0.03, 0.06, 0.06, 0.06, 0.09, 0.175] - RecoveredPerInfectedNoSymptoms = [0.25, 0.25, 0.2, 0.2, 0.2, 0.2] - SeverePerInfectedSymptoms = [0.0075, 0.0075, 0.019, 0.0615, 0.165, 0.225] - CriticalPerSevere = [0.075, 0.075, 0.075, 0.15, 0.3, 0.4] - DeathsPerCritical = [0.05, 0.05, 0.14, 0.14, 0.4, 0.6] - - TimeInfectedNoSymptoms = [2.74, 2.74, 2.565, 2.565, 2.565, 2.565] - TimeInfectedSymptoms = [7.02625, 7.02625, - 7.0665, 6.9385, 6.835, 6.775] - TimeInfectedSevere = [5, 5, 5.925, 7.55, 8.5, 11] - TimeInfectedCritical = [6.95, 6.95, 6.86, 17.36, 17.1, 11.6] - - for i, rho, muCR, muHI, muUH, muDU, tc, ti, th, tu in zip( - range(num_groups), - TransmissionProbabilityOnContact, RecoveredPerInfectedNoSymptoms, - SeverePerInfectedSymptoms, CriticalPerSevere, DeathsPerCritical, - TimeInfectedNoSymptoms, TimeInfectedSymptoms, TimeInfectedSevere, - TimeInfectedCritical): - # Compartment transition duration - model.parameters.TimeExposed[AgeGroup(i)] = 3.335 - model.parameters.TimeInfectedNoSymptoms[AgeGroup(i)] = tc - model.parameters.TimeInfectedSymptoms[AgeGroup(i)] = ti - model.parameters.TimeInfectedSevere[AgeGroup(i)] = th - model.parameters.TimeInfectedCritical[AgeGroup(i)] = tu - - # Compartment transition probabilities - model.parameters.RelativeTransmissionNoSymptoms[AgeGroup(i)] = 1 - model.parameters.TransmissionProbabilityOnContact[AgeGroup(i)] = rho - model.parameters.RecoveredPerInfectedNoSymptoms[AgeGroup(i)] = muCR - model.parameters.RiskOfInfectionFromSymptomatic[AgeGroup(i)] = 0.25 - model.parameters.SeverePerInfectedSymptoms[AgeGroup(i)] = muHI - model.parameters.CriticalPerSevere[AgeGroup(i)] = muUH - model.parameters.DeathsPerCritical[AgeGroup(i)] = muDU - # twice the value of RiskOfInfectionFromSymptomatic - model.parameters.MaxRiskOfInfectionFromSymptomatic[AgeGroup(i)] = 0.5 - # StartDay is the n-th day of the year +def set_covid_parameters( + model, start_date, num_groups=DEFAULT_NUM_AGE_GROUPS): + """Sets COVID-19 specific parameters for all age groups. + + Parameters are based on Kühn et al. (2021): https://doi.org/10.1016/j.mbs.2021.108648 + The function sets age-stratified parameters (when available) the following age groups: + 0-4, 5-14, 15-34, 35-59, 60-79, 80+ + + :param model: MEmilio ODE SECIR model to configure. + :param start_date: Start date of the simulation (used to set StartDay parameter). + :param num_groups: Number of age groups (default: 6). + + """ + # Age-specific transmission and progression parameters + transmission_probability = [0.03, 0.06, 0.06, 0.06, 0.09, 0.175] + recovered_per_infected_no_symptoms = [0.25, 0.25, 0.2, 0.2, 0.2, 0.2] + severe_per_infected_symptoms = [ + 0.0075, 0.0075, 0.019, 0.0615, 0.165, 0.225] + critical_per_severe = [0.075, 0.075, 0.075, 0.15, 0.3, 0.4] + deaths_per_critical = [0.05, 0.05, 0.14, 0.14, 0.4, 0.6] + + # Age-specific compartment transition times (in days) + time_infected_no_symptoms = [2.74, 2.74, 2.565, 2.565, 2.565, 2.565] + time_infected_symptoms = [7.02625, 7.02625, 7.0665, 6.9385, 6.835, 6.775] + time_infected_severe = [5, 5, 5.925, 7.55, 8.5, 11] + time_infected_critical = [6.95, 6.95, 6.86, 17.36, 17.1, 11.6] + + # Apply parameters for each age group + for i in range(num_groups): + age_group = AgeGroup(i) + + model.parameters.TimeExposed[age_group] = 3.335 + model.parameters.TimeInfectedNoSymptoms[age_group] = time_infected_no_symptoms[i] + model.parameters.TimeInfectedSymptoms[age_group] = time_infected_symptoms[i] + model.parameters.TimeInfectedSevere[age_group] = time_infected_severe[i] + model.parameters.TimeInfectedCritical[age_group] = time_infected_critical[i] + + model.parameters.RelativeTransmissionNoSymptoms[age_group] = 1.0 + model.parameters.TransmissionProbabilityOnContact[age_group] = transmission_probability[i] + model.parameters.RiskOfInfectionFromSymptomatic[age_group] = 0.25 + model.parameters.MaxRiskOfInfectionFromSymptomatic[age_group] = 0.5 + model.parameters.RecoveredPerInfectedNoSymptoms[age_group] = recovered_per_infected_no_symptoms[i] + model.parameters.SeverePerInfectedSymptoms[age_group] = severe_per_infected_symptoms[i] + model.parameters.CriticalPerSevere[age_group] = critical_per_severe[i] + model.parameters.DeathsPerCritical[age_group] = deaths_per_critical[i] + + # Set simulation start day model.parameters.StartDay = start_date.day_in_year -def set_contact_matrices(model, data_dir, num_groups=6): - """Setting the contact matrices for a model. +def set_contact_matrices(model, data_dir, num_groups=DEFAULT_NUM_AGE_GROUPS): + """Loads and configures contact matrices for different location types. - :param model: memilio ODE-model, whose contact matrices should be modified. - :param data_dir: directory, where the contact data is stored (should contain folder "contacts"). - :param num_groups: Number of age groups considered. + Contact matrices are loaded for the locations defined in CONTACT_LOCATIONS. - """ + :param model: MEmilio ODE SECIR model to configure. + :param data_dir: Root directory containing contact matrix data (should contain Germany/contacts/ subdirectory). + :param num_groups: Number of age groups (default: 6). + """ contact_matrices = mio.ContactMatrixGroup( - len(list(Location)), num_groups) - locations = ["home", "school_pf_eig", "work", "other"] + len(CONTACT_LOCATIONS), num_groups) - # Loading contact matrices for each location from .txt file - for i, location in enumerate(locations): - baseline_file = os.path.join( - data_dir, "Germany", "contacts", "baseline_" + location + ".txt") + # Load contact matrices for each location + for location_idx, location_name in enumerate(CONTACT_LOCATIONS): + contact_file = os.path.join( + data_dir, "Germany", "contacts", f"baseline_{location_name}.txt" + ) + + if not os.path.exists(contact_file): + raise FileNotFoundError( + f"Contact matrix file not found: {contact_file}" + ) - contact_matrices[i] = mio.ContactMatrix( - mio.read_mobility_plain(baseline_file), + contact_matrices[location_idx] = mio.ContactMatrix( + mio.read_mobility_plain(contact_file) ) + model.parameters.ContactPatterns.cont_freq_mat = contact_matrices -def get_graph( - num_groups, data_dir, mobility_directory, start_date=mio.Date( - 2020, 6, 1), end_date=mio.Date(2020, 8, 31)): - """ Generate the associated graph given the mobility data. +def get_graph(num_groups, data_dir, mobility_directory, start_date, end_date): + """Creates a graph with mobility connections. + + Creates a graph where each node represents a geographic region (here county) with its own ODE model, + and edges represent mobility/commuter connections between these regions. Each node is initialized with + population data and COVID parameters. Edges are weighted by commuter mobility patterns. + + :param num_groups: Number of age groups to model. + :param data_dir: Root directory containing population and contact data. + :param mobility_directory: Path to mobility/commuter data file. + :param start_date: Simulation start date. + :param end_date: Simulation end date (used for data loading). + :returns: Configured ModelGraph with nodes for each region and mobility edges. - :param num_groups: Number of age groups. - :param data_dir: Directory, where the contact data is stored (should contain folder "contacts"). - :param mobility_directory: Directory containing the mobility data. - :param start_date: Date when the simulation starts. - :param end_date: Date when the simulation ends. """ - # Generating and Initializing the model - model = Model(num_groups) - set_covid_parameters(model) - set_contact_matrices(model, data_dir) + model = osecir.Model(num_groups) + set_covid_parameters(model, start_date, num_groups) + set_contact_matrices(model, data_dir, num_groups) - # Generating the graph + # Initialize empty graph graph = osecir.ModelGraph() - # Setting the parameters - scaling_factor_infected = [2.5, 2.5, 2.5, 2.5, 2.5, 2.5] + # To account for underreporting + scaling_factor_infected = [2.5] * num_groups scaling_factor_icu = 1.0 - tnt_capacity_factor = 7.5 / 100000. + # Test & trace capacity as in Kühn et al. (2021): https://doi.org/10.1016/j.mbs.2021.108648 + tnt_capacity_factor = 7.5 / 100000.0 - # Path containing the population data - data_dir_Germany = os.path.join(data_dir, "Germany") - pydata_dir = os.path.join(data_dir_Germany, "pydata") + # Paths to population data + pydata_dir = os.path.join(data_dir, "Germany", "pydata") + path_population_data = os.path.join( + pydata_dir, "county_current_population.json") - path_population_data = os.path.join(pydata_dir, - "county_current_population.json") + # Verify data files exist + if not os.path.exists(path_population_data): + raise FileNotFoundError( + f"Population data not found: {path_population_data}" + ) + # Create one node per county is_node_for_county = True - mio.osecir.set_nodes( - model.parameters, mio.Date( - start_date.year, start_date.month, start_date.day), - mio.Date(end_date.year, end_date.month, end_date.day), - pydata_dir, path_population_data, is_node_for_county, graph, - scaling_factor_infected, scaling_factor_icu, tnt_capacity_factor, 0, False) - - mio.osecir.set_edges( - mobility_directory, graph, len(Location)) + # Populate graph nodes with county level data + osecir.set_nodes( + model.parameters, + start_date, + end_date, + pydata_dir, + path_population_data, + is_node_for_county, + graph, + scaling_factor_infected, + scaling_factor_icu, + tnt_capacity_factor, + 0, # Zero days extrapolating the data + False # No extrapolation of data + ) + + # Add mobility edges between regions + osecir.set_edges(mobility_directory, graph, len(CONTACT_LOCATIONS)) return graph +def get_compartment_factors(): + """Draw base factors for compartment initialization. + + Factors follow the sampling strategy from Schmidt et al. + (2024): symptomatic individuals between 0.01% and 5% of the population, + exposed and asymptomatic compartments proportional to that symptomatic + proportion, and hospital/ICU/deaths sampled hierarchically. + Recovered is taken from the remaining feasible proportion and susceptibles + fill the residual. + """ + p_infected = random.uniform(0.0001, 0.05) + p_exposed = p_infected * random.uniform(0.1, 5.0) + p_ins = p_infected * random.uniform(0.1, 5.0) + p_hosp = p_infected * random.uniform(0.001, 1.0) + p_icu = p_hosp * random.uniform(0.001, 1.0) + p_dead = p_icu * random.uniform(0.001, 1.0) + + sum_randoms = ( + p_infected + p_exposed + p_ins + p_hosp + p_icu + p_dead + ) + + if sum_randoms >= 1.0: + raise RuntimeError( + "Sampled compartment factors exceed total population. Adjust bounds." + ) + + p_recovered = random.uniform(0.0, 1.0 - sum_randoms) + p_susceptible = max(1.0 - (sum_randoms + p_recovered), 0.0) + + return { + "infected": p_infected, + "exposed": p_exposed, + "infected_no_symptoms": p_ins, + "hospitalized": p_hosp, + "critical": p_icu, + "dead": p_dead, + "recovered": p_recovered, + "susceptible": p_susceptible + } + + +def _initialize_compartments_for_node( + model, factors, num_groups, within_group_variation): + """Initializes epidemic compartments using shared base factors. + + :param model: Model instance for the specific node. + :param factors: Compartment factors obtained from get_compartment_factors. + :param num_groups: Number of age groups. + :param within_group_variation: Whether to apply additional random scaling per age group. + """ + + def _variation(): + return random.uniform(0.1, 1.0) if within_group_variation else 1.0 + + for age_idx in range(num_groups): + age_group = AgeGroup(age_idx) + total_population = model.populations.get_group_total_AgeGroup( + age_group) + + infected_symptoms = total_population * \ + factors["infected"] * _variation() + + exposed = infected_symptoms * factors["exposed"] * _variation() + infected_no_symptoms = infected_symptoms * \ + factors["infected_no_symptoms"] * _variation() + infected_severe = infected_symptoms * \ + factors["hospitalized"] * _variation() + infected_critical = infected_severe * \ + factors["critical"] * _variation() + dead = infected_critical * factors["dead"] * _variation() + recovered = total_population * factors["recovered"] * _variation() + + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedSymptoms)] = infected_symptoms + model.populations[age_group, Index_InfectionState( + InfectionState.Exposed)] = exposed + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedNoSymptoms)] = infected_no_symptoms + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedSevere)] = infected_severe + model.populations[age_group, Index_InfectionState( + InfectionState.InfectedCritical)] = infected_critical + model.populations[age_group, Index_InfectionState( + InfectionState.Dead)] = dead + model.populations[age_group, Index_InfectionState( + InfectionState.Recovered)] = recovered + + # Susceptibles are the remainder + model.populations.set_difference_from_group_total_AgeGroup( + (age_group, InfectionState.Susceptible), total_population + ) + + +def _apply_dampings_to_model(model, damping_days, damping_factors, num_groups): + """Applies contact dampings (NPIs) to model at specified days. + + Currently only supports global dampings (same for all age groups and spatial units). + + :param model: Model to apply dampings to. + :param damping_days: Days at which to apply dampings. + :param damping_factors: Multiplicative factors for contact reduction. + :param num_groups: Number of age groups. + :returns: Tuple of (damped_contact_matrices, damping_coefficients). + + """ + damped_matrices = [] + damping_coefficients = [] + + for day, factor in zip(damping_days, damping_factors): + # Create uniform damping matrix for all age groups + damping_matrix = np.ones((num_groups, num_groups)) * factor + + # Add damping to model + model.parameters.ContactPatterns.cont_freq_mat.add_damping( + Damping(coeffs=damping_matrix, t=day, level=0, type=0) + ) + + # Store resulting contact matrix and coefficients + damped_matrices.append( + model.parameters.ContactPatterns.cont_freq_mat.get_matrix_at( + day + 1)) + damping_coefficients.append(damping_matrix) + + return damped_matrices, damping_coefficients + + def run_secir_groups_simulation( - days, damping_days, damping_factors, graph, num_groups=6): - """Run an ODE SECIR simulation with multiple age groups and regions. - - Uses an ODE SECIR model allowing for asymptomatic infection with 6 - different age groups. The model is stratified by region using a graph structure. - Virus-specific parameters are fixed and initial number of persons in the - particular infection states are chosen randomly between predefined bounds. - - :param days: Number of days simulated within a single run. - :param damping_days: List of days where dampings are applied. - :param damping_factors: List of damping factors associated to the damping days. - :param graph: Graph initialized for the start_date with the population data. - :param num_groups: Number of age groups considered in the simulation. - :param start_date: Date when the simulation starts. - :returns: Tuple containing (simulation_results, damped_matrices, damping_coefficients, runtime) + days, damping_days, damping_factors, graph, within_group_variation, + num_groups=DEFAULT_NUM_AGE_GROUPS): + """Runs a multi-region ODE SECIR simulation with age groups and NPIs. + + Performs the following steps: + 1. Initialize each node with compartment specific (random) factors + 2. Apply contact dampings (NPIs) at specified days + 3. Run the simulation + 4. Post-process and return results + + :param days: Total number of days to simulate. + :param damping_days: List of days when NPIs are applied. + :param damping_factors: List of contact reduction factors for each damping. + :param graph: Pre-configured ModelGraph with nodes and edges. + :param within_group_variation: Whether to apply per-age random variation during initialization. + :param num_groups: Number of age groups (default: 6). + :returns: Tuple containing dataset_entry (simulation results for each node [num_nodes, time_steps, compartments]), + damped_matrices (contact matrices at each damping day), damping_coefficients (damping coefficient matrices), + and runtime (execution time in seconds). + :raises ValueError: If lengths of damping_days and damping_factors don't match. + """ if len(damping_days) != len(damping_factors): - raise ValueError("Length of damping_days and damping_factors differ!") + raise ValueError( + f"Length mismatch: damping_days ({len(damping_days)}) != " + f"damping_factors ({len(damping_factors)})" + ) + + # Initialize each node in the graph + damped_matrices = [] + damping_coefficients = [] + + factors = get_compartment_factors() + + for node_idx in range(graph.num_nodes): + model = graph.get_node(node_idx).property + + # Initialize compartment populations + _initialize_compartments_for_node( + model, factors, num_groups, within_group_variation) + + # Apply dampings/NPIs + if damping_days: + node_damped_mats, node_damping_coeffs = _apply_dampings_to_model( + model, damping_days, damping_factors, num_groups + ) + # Store only from first node, since dampings are global at the moment + if node_idx == 0: + damped_matrices = node_damped_mats + damping_coefficients = node_damping_coeffs - # Load ground truth bounds for initialization - pydata_dir = os.path.join(data_dir, "Germany", "pydata") - upper_bound_dir = os.path.join( - pydata_dir, "ground_truth_upper_bound.pickle") - lower_bound_dir = os.path.join( - pydata_dir, "ground_truth_lower_bound.pickle") - with open(upper_bound_dir, 'rb') as f: - ground_truth_upper_bound = pickle.load(f) - with open(lower_bound_dir, 'rb') as f: - ground_truth_lower_bound = pickle.load(f) - - # Initialize model for each node - for node_indx in range(graph.num_nodes): - model = graph.get_node(node_indx).property - - # Iterate over the different age groups - for i in range(num_groups): - age_group = AgeGroup(i) - pop_age_group = model.populations.get_group_total_AgeGroup( - age_group) - - # Generate initial compartment populations with random values between bounds - valid_configuration = False - while valid_configuration is False: - init_data = np.asarray([0 for _ in range(8)]) - max_data = ground_truth_upper_bound[node_indx][:, i] - min_data = ground_truth_lower_bound[node_indx][:, i] - for j in range(1, 8): - init_data[j] = random.uniform(min_data[j], max_data[j]) - - # Check if configuration is valid (infected population < total population) - if np.sum(init_data[1:]) < pop_age_group: - valid_configuration = True - - # Set the populations for the different compartments - model.populations[age_group, Index_InfectionState( - InfectionState.Exposed)] = init_data[1] - model.populations[age_group, Index_InfectionState( - InfectionState.InfectedNoSymptoms)] = init_data[2] - model.populations[age_group, Index_InfectionState( - InfectionState.InfectedSymptoms)] = init_data[3] - model.populations[age_group, Index_InfectionState( - InfectionState.InfectedSevere)] = init_data[4] - model.populations[age_group, Index_InfectionState( - InfectionState.InfectedCritical)] = init_data[5] - model.populations[age_group, Index_InfectionState( - InfectionState.Recovered)] = init_data[6] - model.populations[age_group, Index_InfectionState( - InfectionState.Dead)] = init_data[7] - model.populations.set_difference_from_group_total_AgeGroup(( - age_group, InfectionState.Susceptible), pop_age_group) - - # Apply damping information (currently only global dampings supported) - damped_matrices = [] - damping_coefficients = [] - - for i in np.arange(len(damping_days)): - day = damping_days[i] - factor = damping_factors[i] - - damping = np.ones((num_groups, num_groups)) * np.float16(factor) - model.parameters.ContactPatterns.cont_freq_mat.add_damping(Damping( - coeffs=(damping), t=day, level=0, type=0)) - damped_matrices.append( - model.parameters.ContactPatterns.cont_freq_mat.get_matrix_at( - day + 1)) - damping_coefficients.append(damping) - - # Apply mathematical constraints to parameters model.apply_constraints() - # Set model to graph - graph.get_node(node_indx).property.populations = model.populations + # Update graph with initialized populations + graph.get_node(node_idx).property.populations = model.populations + + # Run simulation and measure runtime + study = ParameterStudy(graph, t0=0, tmax=days, dt=0.5, num_runs=1) - # Run simulation - study = ParameterStudy(graph, 0, days, dt=0.5, num_runs=1) start_time = time.perf_counter() - study.run() + study_results = study.run() runtime = time.perf_counter() - start_time - # Process results - graph_run = study.run()[0] + # Interpolate results to daily values + graph_run = study_results[0] results = interpolate_simulation_result(graph_run) - for result_indx in range(len(results)): - results[result_indx] = remove_confirmed_compartments( - np.asarray(results[result_indx]), num_groups) + # Remove confirmed compartments (not used in GNN) + for result_idx in range(len(results)): + results[result_idx] = remove_confirmed_compartments( + np.asarray(results[result_idx]), num_groups + ) dataset_entry = copy.deepcopy(results) return dataset_entry, damped_matrices, damping_coefficients, runtime -def generate_data( - num_runs, data_dir, path, input_width, label_width, start_date, - end_date, save_data=True, transform=True, damping_method="classic", - max_number_damping=3, mobility_file="commuter_mobility_2022.txt"): - """Generate dataset by calling run_secir_groups_simulation multiple times. +def generate_data(num_runs, data_dir, output_path, input_width, label_width, + start_date, end_date, save_data=True, transform=True, + damping_method="classic", max_number_damping=3, + mobility_file="commuter_mobility.txt", + num_groups=DEFAULT_NUM_AGE_GROUPS, + within_group_variation=True): + """Generates training dataset for GNN surrogate model. + + Runs num_runs-ODE SECIR simulations with random initial conditions and damping patterns to create a training dataset. + If save_data=True, saves a pickle file named: + 'GNN_data_{label_width}days_{max_number_damping}dampings_{damping_method}{num_runs}.pickle' :param num_runs: Number of simulation runs to generate. - :param data_dir: Directory with all data needed to initialize the models. - :param path: Path where the datasets are stored. - :param input_width: Number of time steps used for model input. - :param label_width: Number of time steps (days) used as model output/label. - :param start_date: Date when the simulation starts. - :param end_date: Date when the simulation ends. - :param save_data: Option to deactivate the save of the dataset. Default True. - :param transform: Option to deactivate the transformation of the data. Default True. - :param damping_method: String specifying the damping method. Values: "classic", "active", "random". - :param max_number_damping: Maximal number of possible dampings. - :param mobility_file: Filename of the mobility data file. - :returns: Dictionary containing inputs, labels, contact_matrix, damping_days and damping_factors. + :param data_dir: Root directory containing all required input data (population, contacts, mobility). + :param output_path: Directory where generated dataset will be saved. + :param input_width: Number of time steps for model input. + :param label_width: Number of time steps for model output/labels. + :param start_date: Simulation start date. + :param end_date: Simulation end date (used for set_nodes function). + :param save_data: Whether to save dataset (default: True). + :param transform: Whether to apply scaling transformation to data (default: True). + :param damping_method: Method for generating damping patterns: "classic", "active", "random". + :param max_number_damping: Maximum number of damping events per simulation. + :param mobility_file: Filename of mobility file (in data_dir/Germany/mobility/). + :param num_groups: Number of age groups (default: 6). + :param within_group_variation: Whether to apply random variation per spatial unit/age group when initializing compartments. + :returns: Dictionary with keys: "inputs" ([num_runs, input_width, num_nodes, features]), + "labels" ([num_runs, label_width, num_nodes, features]), + "contact_matrix" (List of damped contact matrices), "damping_days" (List of damping day arrays), + "damping_factors" (List of damping factor arrays). + """ set_log_level(mio.LogLevel.Error) - days = label_width + input_width - 1 - # Preparing output dictionary + # Calculate total simulation days + total_days = label_width + input_width - 1 + + # Initialize output dictionary data = { "inputs": [], "labels": [], @@ -313,95 +460,163 @@ def generate_data( "damping_factors": [] } - # Setting basic parameter - num_groups = 6 - mobility_dir = data_dir + "/Germany/mobility/" + mobility_file - graph = get_graph(num_groups, data_dir, mobility_dir, start_date, end_date) + # Build mobility file path + mobility_path = os.path.join( + data_dir, "Germany", "mobility", mobility_file) + + # Verify mobility file exists + if not os.path.exists(mobility_path): + raise FileNotFoundError(f"Mobility file not found: {mobility_path}") + + # Create graph (reused for all runs with different initial conditions) + graph = get_graph(num_groups, data_dir, + mobility_path, start_date, end_date) - # show progess in terminal for longer runs - # Due to the random structure, there is currently no need to shuffle the data - bar = Bar('Number of Runs done', max=num_runs) + print(f"\nGenerating {num_runs} simulation runs...") + bar = Bar( + 'Progress', max=num_runs, + suffix='%(percent)d%% [%(elapsed_td)s / %(eta_td)s]') - times = [] - for i in range(0, num_runs): - # Generate random damping days and damping factors + runtimes = [] + + for _ in range(num_runs): + # Generate random damping pattern if max_number_damping > 0: damping_days, damping_factors = dampings.generate_dampings( - days, max_number_damping, method=damping_method, - min_distance=2, min_damping_day=2) + total_days, + max_number_damping, + method=damping_method, + min_distance=2, + min_damping_day=2 + ) else: damping_days = [] damping_factors = [] + # Run simulation - data_run, damped_matrices, damping_coefficients, t_run = run_secir_groups_simulation( - days, damping_days, damping_factors, graph, num_groups) + simulation_result, damped_mats, damping_coeffs, runtime = \ + run_secir_groups_simulation( + total_days, damping_days, damping_factors, graph, + within_group_variation, num_groups + ) - times.append(t_run) + runtimes.append(runtime) - inputs = np.asarray(data_run).transpose(1, 0, 2)[: input_width] - labels = np.asarray(data_run).transpose(1, 0, 2)[input_width:] + # Split into inputs and labels + # Shape: [num_nodes, time_steps, features] -> transpose to [time_steps, num_nodes, features] + result_transposed = np.asarray(simulation_result).transpose(1, 0, 2) + inputs = result_transposed[:input_width] + labels = result_transposed[input_width:] + # Store results data["inputs"].append(inputs) data["labels"].append(labels) - data["contact_matrix"].append(np.array(damped_matrices)) - data["damping_factors"].append(damping_coefficients) + data["contact_matrix"].append(np.array(damped_mats)) + data["damping_factors"].append(damping_coeffs) data["damping_days"].append(damping_days) bar.next() bar.finish() - print( - f"For Days = {days}, AVG runtime: {np.mean(times)}s, Median runtime: {np.median(times)}s") + # Print performance statistics + print(f"\nSimulation Statistics:") + print(f" Total days simulated: {total_days}") + print(f" Average runtime: {np.mean(runtimes):.3f}s") + print(f" Median runtime: {np.median(runtimes):.3f}s") + print(f" Total time: {np.sum(runtimes):.1f}s") + # Save dataset if requested if save_data: + # Apply scaling transformation + inputs_scaled, labels_scaled = scale_data(data, transform) - inputs, labels = scale_data(data, transform) - - all_data = {"inputs": inputs, - "labels": labels, - "damping_day": data["damping_days"], - "contact_matrix": data["contact_matrix"], - "damping_coeff": data["damping_factors"] - } + all_data = { + "inputs": inputs_scaled, + "labels": labels_scaled, + "damping_day": data["damping_days"], + "contact_matrix": data["contact_matrix"], + "damping_coeff": data["damping_factors"] + } - # check if data directory exists. If necessary create it. - if not os.path.isdir(path): - os.mkdir(path) + # Create output directory if needed + os.makedirs(output_path, exist_ok=True) - # generate the filename + # Generate filename if num_runs < 1000: - filename = 'GNN_data_%ddays_%ddampings_' % ( - label_width, max_number_damping) + damping_method+'%d.pickle' % (num_runs) + filename = f'GNN_data_{label_width}days_{max_number_damping}dampings_{damping_method}{num_runs}.pickle' else: - filename = 'GNN_data_%ddays_%ddampings_' % ( - label_width, max_number_damping) + damping_method+'%dk.pickle' % (num_runs//1000) + filename = f'GNN_data_{label_width}days_{max_number_damping}dampings_{damping_method}{num_runs//1000}k.pickle' - # save dict to pickle file - with open(os.path.join(path, filename), 'wb') as f: + # Save to pickle file + output_file = os.path.join(output_path, filename) + with open(output_file, 'wb') as f: pickle.dump(all_data, f) + print(f"\nDataset saved to: {output_file}") + return data -if __name__ == "__main__": +def main(): + """Main function for dataset generation. - path = os.getcwd() - path_output = os.path.join(os.getcwd(), 'saves') - data_dir = os.path.join(os.getcwd(), 'data') - input_width = 5 - number_of_dampings = 0 - num_runs = 1 - label_width_list = [30] + Example configuration for generating GNN training data. + + """ + # Set random seed for reproducibility random.seed(10) - # Define the start and the latest end date for the simulation + # Configuration + data_dir = os.path.join(os.getcwd(), 'data') + output_path = os.path.join(os.getcwd(), 'saves') + + # Simulation parameters + input_width = 5 # Days of history used as input + num_runs = 1 # Number of simulation runs + max_dampings = 0 # Number of NPI dampings per simulation + + # Prediction horizons to generate data for + prediction_horizons = [30] # Days to predict into the future + + # Simulation time period start_date = mio.Date(2020, 10, 1) end_date = mio.Date(2021, 10, 31) - for label_width in label_width_list: + # Generate datasets + print("=" * 70) + print("GNN Surrogate Model - Dataset Generation") + print("=" * 70) + print(f"Data directory: {data_dir}") + print(f"Output directory: {output_path}") + print(f"Simulation period: {start_date} to {end_date}") + print(f"Number of runs per configuration: {num_runs}") + print(f"Input width: {input_width} days") + print(f"Max dampings: {max_dampings}") + print("=" * 70) + + for label_width in prediction_horizons: + print(f"\n{'='*70}") + print(f"Generating data for {label_width}-day predictions") + print(f"{'='*70}") + generate_data( - num_runs=num_runs, data_dir=data_dir, path=path_output, - input_width=input_width, label_width=label_width, - start_date=start_date, end_date=end_date, save_data=True, - damping_method="active", max_number_damping=number_of_dampings) + num_runs=num_runs, + data_dir=data_dir, + output_path=output_path, + input_width=input_width, + label_width=label_width, + start_date=start_date, + end_date=end_date, + save_data=True, + damping_method="active", + max_number_damping=max_dampings + ) + + print(f"\n{'='*70}") + print("Dataset generation complete!") + print(f"{'='*70}") + + +if __name__ == "__main__": + main() From 304a702f7f11714292ca0e50a865ccc4782348bc Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 11:57:41 +0100 Subject: [PATCH 25/41] . --- .../memilio/surrogatemodel/GNN/data_generation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index c91a58849d..eb3e44e271 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -30,6 +30,7 @@ import pickle import random import time +from enum import Enum from typing import Dict, List, Tuple import numpy as np @@ -569,7 +570,7 @@ def main(): # Configuration data_dir = os.path.join(os.getcwd(), 'data') - output_path = os.path.join(os.getcwd(), 'saves') + output_path = os.path.join(os.getcwd(), 'generated_datasets') # Simulation parameters input_width = 5 # Days of history used as input From d19c6abf8b098af3960e19afbaa60cd6f056968e Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 12:24:16 +0100 Subject: [PATCH 26/41] [ci skip] rework evaluate and trian --- .../surrogatemodel/GNN/evaluate_and_train.py | 603 ++++++++++-------- 1 file changed, 324 insertions(+), 279 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index d9cce1fc5a..6517030f37 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -17,13 +17,21 @@ # See the License for the specific language governing permissions and # limitations under the License. ############################################################################# +""" +Training and evaluation module for GNN-based surrogate models. + +This module loads GNN training data, trains/evaluates surrogate models, and saves weights plus metrics. +""" -import os import pickle -import spektral import time -import pandas as pd +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Dict, Iterable, List, Tuple + import numpy as np +import pandas as pd +import spektral from tensorflow.keras.optimizers import Adam @@ -35,298 +43,328 @@ from memilio.surrogatemodel.utils.helper_functions import (calc_split_index) from spektral.data import MixedLoader -from spektral.layers import ( - GCSConv, GlobalAvgPool, GlobalAttentionPool, ARMAConv, AGNNConv, APPNPConv, - CrystalConv, GATConv, GINConv, XENetConv, GCNConv, GCSConv) +from spektral.layers import ARMAConv from spektral.utils.convolution import normalized_laplacian, rescale_laplacian -def create_dataset(path_cases, path_mobility, number_of_nodes=400): - """ - Create a dataset from the given file path. - :param path_cases: Path to the pickle file containing case data. - :param path_mobility: Path to the directory containing mobility data. - .param number_of_nodes: Number of nodes in the graph. - :return: A Spektral Dataset object containing the processed data. +def _iterate_batches(loader): + """Yield all batches from a Spektral loader for one epoch.""" + for _ in range(loader.steps_per_epoch): + yield next(loader) + + +class StaticGraphDataset(spektral.data.Dataset): + """Spektral dataset wrapper for samples that share a single adjacency matrix.""" + + def __init__( + self, + node_features, + node_labels, + adjacency): + self._node_features = node_features + self._node_labels = node_labels + self._adjacency = adjacency.astype(np.float32) + super().__init__() + + def read(self): + """Create one Graph object per sample.""" + return [ + spektral.data.Graph( + x=feature.astype(np.float32), + y=label.astype(np.float32), + a=self._adjacency, + ) + for feature, label in zip(self._node_features, self._node_labels) + ] + + +@dataclass +class TrainingSummary: + """Container for aggregated training and evaluation metrics.""" + + model_name: str + mean_train_loss: float + mean_val_loss: float + mean_test_loss: float + mean_test_loss_orig: float + training_minutes: float + train_losses: List[List[float]] + val_losses: List[List[float]] + + +def load_gnn_dataset( + dataset_path, + mobility_dir, + number_of_nodes=400, + mobility_filename="commuter_mobility_2022.txt"): + """Load serialized samples and mobility data into a Spektral dataset. + + Args: + dataset_path: Pickle file containing the generated training samples. + mobility_dir: Directory containing the commuter mobility file. + number_of_nodes: Number of spatial nodes to retain from the mobility data. + mobility_filename: Mobility file name to use. + + Returns: + Spektral dataset with one `Graph` per sample sharing a common adjacency matrix. """ + dataset_path = Path(dataset_path) + mobility_dir = Path(mobility_dir) + + if not dataset_path.exists(): + raise FileNotFoundError(f"Dataset not found: {dataset_path}") - # Load data from pickle file - file = open( - path_cases, 'rb') - data = pickle.load(file) - # Extract inputs and labels from the loaded data - inputs = data['inputs'] - labels = data['labels'] + mobility_path = mobility_dir / mobility_filename + if not mobility_path.exists(): + raise FileNotFoundError(f"Mobility file not found: {mobility_path}") - len_dataset = len(inputs) + with dataset_path.open("rb") as fp: + data = pickle.load(fp) - if len_dataset == 0: # check if dataset is empty + if "inputs" not in data or "labels" not in data: + raise KeyError( + f"Dataset at {dataset_path} must contain 'inputs' and 'labels' keys." + ) + + inputs = np.asarray(data["inputs"]) + labels = np.asarray(data["labels"]) + if inputs.shape[0] == 0: raise ValueError( - "Dataset is empty. Please provide a valid dataset with at least one sample.") - # Calculate the flattened shape of inputs and labels - shape_input_flat = np.asarray( - inputs).shape[1]*np.asarray(inputs).shape[2] - shape_labels_flat = np.asarray( - labels).shape[1]*np.asarray(labels).shape[2] - # Reshape inputs and labels to the required format - new_inputs = np.asarray( - inputs.transpose(0, 3, 1, 2)).reshape( - len_dataset, number_of_nodes, shape_input_flat) - new_labels = np.asarray(labels.transpose(0, 3, 1, 2)).reshape( - len_dataset, number_of_nodes, shape_labels_flat) - - # Load mobility data and create adjacency matrix - commuter_file = open(os.path.join( - path_mobility, 'commuter_mobility_2022.txt'), 'rb') - commuter_data = pd.read_csv(commuter_file, sep=" ", header=None) - sub_matrix = commuter_data.iloc[:number_of_nodes, 0:number_of_nodes] - - adjacency_matrix = np.asarray(sub_matrix) - - adjacency_matrix[adjacency_matrix > 0] = 1 # make the adjacency binary - - node_features = new_inputs - - node_labels = new_labels - - # Define a custom Dataset class - class MyDataset(spektral.data.dataset.Dataset): - def read(self): - self.a = adjacency_matrix - return [spektral.data.Graph(x=x, y=y) for x, y in zip(node_features, node_labels)] - - super().__init__(**kwargs) - # Instantiate the custom dataset - data = MyDataset() - - return data - - -def train_step(inputs, target, loss_fn, model, optimizer): - '''Perform a single training step. - :param inputs: Tuple (x, a) where x is the node features and a is the adjacency matrix. - :param target: Ground truth labels. - :param loss_fn: Loss function to use. - :param model: The GNN model to train. - :param optimizer: Optimizer to use for training. - :return: Loss and accuracy for the training step.''' - # Record operations for automatic differentiation + "Loaded dataset is empty; expected at least one sample.") + + # Flatten temporal dimensions into feature vectors per node. + num_samples, input_width, num_nodes, num_features = inputs.shape + _, label_width, _, label_features = labels.shape + + if num_nodes != number_of_nodes: + raise ValueError( + f"Number of nodes in dataset ({num_nodes}) does not match expected " + f"value ({number_of_nodes}).") + + node_features = inputs.transpose(0, 2, 1, 3).reshape( + num_samples, number_of_nodes, input_width * num_features) + node_labels = labels.transpose(0, 2, 1, 3).reshape( + num_samples, number_of_nodes, label_width * label_features) + + commuter_data = pd.read_csv(mobility_path, sep=" ", header=None) + adjacency_matrix = commuter_data.iloc[ + :number_of_nodes, :number_of_nodes + ].to_numpy() + adjacency_matrix = (adjacency_matrix > 0).astype(np.float32) + adjacency_matrix = np.maximum(adjacency_matrix, adjacency_matrix.T) + + return StaticGraphDataset(node_features, node_labels, adjacency_matrix) + + +def create_dataset(path_cases, path_mobility, number_of_nodes=400): + """Compatibility wrapper around `load_gnn_dataset`.""" + return load_gnn_dataset( + Path(path_cases), + Path(path_mobility), + number_of_nodes=number_of_nodes) + + +def _train_step_impl(model, optimizer, loss_fn, inputs, target): + """Perform one optimization step.""" with tf.GradientTape() as tape: predictions = model(inputs, training=True) - loss = loss_fn(target, predictions) + \ - sum(model.losses) # Add regularization losses - # Compute gradients and update model weights + loss = loss_fn(target, predictions) + if model.losses: + loss += tf.add_n(model.losses) + gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) - acc = tf.reduce_mean(loss_fn(target, predictions)) - return loss, acc + + metric = tf.reduce_mean(loss_fn(target, predictions)) + return loss, metric + + +def train_step(*args, **kwargs): + """Wrapper that accepts both old and new call signatures.""" + if len(args) == 5 and not kwargs: + first, second, third, fourth, fifth = args + if isinstance(first, (tuple, list)): + inputs, target = first, second + loss_fn = third + model = fourth + optimizer = fifth + else: + model, optimizer, loss_fn, inputs, target = args + else: + raise TypeError( + "train_step expects either (inputs, target, loss_fn, model, optimizer) " + "or (model, optimizer, loss_fn, inputs, target).") + + return _train_step_impl(model, optimizer, loss_fn, inputs, target) def evaluate(loader, model, loss_fn, retransform=False): - '''Evaluate the model on a validation or test set. - :param loader: Data loader for the evaluation set. - :param model: The GNN model to evaluate. - :param loss_fn: Loss function to use. - :param retransform: Whether to apply inverse transformation to the outputs. - :return: Losses and mean_loss for the evaluation set.''' - output = [] - step = 0 - while step < loader.steps_per_epoch: - step += 1 - inputs, target = loader.__next__() - pred = model(inputs, training=False) + """Evaluate the model on the dataset provided by loader.""" + total_loss = 0.0 + total_metric = 0.0 + total_samples = 0 + + for inputs, target in _iterate_batches(loader): + predictions = model(inputs, training=False) + + target_tensor = tf.convert_to_tensor(target, dtype=tf.float32) + prediction_tensor = tf.cast(predictions, tf.float32) if retransform: - target = np.expm1(target) - pred = np.expm1(pred) - # Calculate loss and metrics - outs = ( - loss_fn(target, pred), - tf.reduce_mean(loss_fn(target, pred)), - len(target), # Keep track of batch size - ) - output.append(outs) - # Aggregate results at the end of the epoch - if step == loader.steps_per_epoch: - output = np.array(output) - return np.average(output[:, :-1], 0, weights=output[:, -1]) + target_tensor = tf.math.expm1(target_tensor) + prediction_tensor = tf.math.expm1(prediction_tensor) + + batch_losses = loss_fn(target_tensor, prediction_tensor) + batch_loss = tf.reduce_mean(batch_losses) + batch_metric = tf.reduce_mean(batch_losses) + + batch_size_tensor = tf.cast(tf.shape(target_tensor)[0], tf.float32) + batch_size = float(batch_size_tensor.numpy()) + total_loss += float(batch_loss.numpy()) * batch_size + total_metric += float(batch_metric.numpy()) * batch_size + total_samples += int(round(batch_size)) + + if total_samples == 0: + return 0.0, 0.0 + + mean_loss = total_loss / total_samples + mean_metric = total_metric / total_samples + return mean_loss, mean_metric def train_and_evaluate( data, batch_size, epochs, model, loss_fn, optimizer, es_patience, - save_dir="", save_name="model"): - '''Train and evaluate the GNN model. - :param data: The dataset to use for training and evaluation. - :param batch_size: Batch size for training. - :param epochs: Maximum number of epochs to train. - :param model: The GNN model to train and evaluate. - :param loss_fn: Loss function to use. - :param optimizer: Optimizer to use for training. - :param es_patience: Patience for early stopping. - :param save_dir: Directory to save results into. If not provided, results are not saved. - :param save_name: Name to use when saving the results (without extension). - :return: A dictionary containing training and evaluation results if save_results is False.''' - n = len(data) - # Determine split indices for training, validation, and test sets + save_dir=None, save_name="model"): + """Train the provided GNN model.""" + dataset_size = len(data) + if dataset_size == 0: + raise ValueError("Dataset must contain at least one sample.") + n_train, n_valid, n_test = calc_split_index( - n, split_train=0.7, split_valid=0.2, split_test=0.1) - - # Split data into train, validation, and test sets - train_data, valid_data, test_data = data[:n_train], data[n_train:n_train + - n_valid], data[n_train+n_valid:] - - # Build the model by calling it on a batch of data - loader = MixedLoader(data) - inputs, _ = loader.__next__() - model(inputs) - - # Define Data Loaders - loader_tr = MixedLoader( - train_data, batch_size=batch_size, epochs=epochs, shuffle=False) - loader_val = MixedLoader( - valid_data, batch_size=batch_size, shuffle=False) - loader_test = MixedLoader( - test_data, batch_size=n_test, shuffle=False) - - df = pd.DataFrame(columns=[ - "train_loss", "val_loss", "test_loss", - "test_loss_orig", "training_time", - "loss_history", "val_loss_history"]) - - # Initialize variables to track best scores and histories - test_scores = [] - test_scores_r = [] - train_losses = [] - val_losses = [] - - losses_history_all = [] - val_losses_history_all = [] + dataset_size, split_train=0.7, split_valid=0.2, split_test=0.1) + if n_train == 0 or n_valid == 0 or n_test == 0: + raise ValueError( + "Dataset split produced empty partitions. Provide a larger dataset " + "or adjust the split configuration.") + + train_data = data[:n_train] + valid_data = data[n_train:n_train + n_valid] + test_data = data[n_train + n_valid:] + + # Build the model by passing a single batch through it. + build_loader = MixedLoader(train_data, batch_size=min( + batch_size, max(1, n_train)), epochs=1, shuffle=False) + build_inputs, _ = next(build_loader) + model(build_inputs) + + def _make_loader(dataset, *, batch_size, shuffle=False): + return MixedLoader( + dataset, batch_size=batch_size, epochs=1, shuffle=shuffle) best_val_loss = np.inf - best_weights = None - patience = es_patience - results = [] - losses_history = [] - val_losses_history = [] - - ################################################################################ - # Train model - ################################################################################ - start = time.perf_counter() - step = 0 - epoch = 0 - for batch in loader_tr: - step += 1 - loss, acc = train_step(*batch, loss_fn, model, optimizer) - results.append((loss, acc)) - - # Compute validation loss and accuracy at the end of each epoch - if step == loader_tr.steps_per_epoch: - step = 0 - epoch += 1 - val_loss, val_acc = evaluate(loader_val, model, loss_fn) - print( - "Ep. {} - Loss: {:.3f} - Acc: {:.3f} - Val loss: {:.3f} - Val acc: {:.3f}".format( - epoch, *np.mean(results, 0), val_loss, val_acc - ) - ) - # Check if loss improved for early stopping - if val_loss < best_val_loss: - best_val_loss = val_loss - patience = es_patience - print(f"New best val_loss {val_loss:.3f}") - best_weights = model.get_weights() - else: - patience -= 1 - if patience == 0: - print( - "Early stopping (best val_loss: {})".format( - best_val_loss)) - break - results = [] - losses_history.append(loss) - val_losses_history.append(val_loss) - ################################################################################ - # Evaluate model - ################################################################################ - # Load best model weights before evaluation on test set + best_weights = model.get_weights() + patience_counter = es_patience + + epoch_train_losses: List[float] = [] + epoch_val_losses: List[float] = [] + + start_time = time.perf_counter() + + for epoch in range(1, epochs + 1): + train_loader = _make_loader( + train_data, batch_size=batch_size, shuffle=True) + batch_losses = [] + for inputs, target in _iterate_batches(train_loader): + loss, _ = train_step(model, optimizer, loss_fn, inputs, target) + batch_losses.append(float(loss.numpy())) + + epoch_train_loss = float(np.mean(batch_losses) + ) if batch_losses else 0.0 + epoch_train_losses.append(epoch_train_loss) + + val_loader = _make_loader(valid_data, batch_size=min( + batch_size, max(1, n_valid)), shuffle=False) + val_loss, _ = evaluate(val_loader, model, loss_fn) + epoch_val_losses.append(val_loss) + + print( + f"Epoch {epoch:02d} | train_loss={epoch_train_loss:.4f} " + f"| val_loss={val_loss:.4f}" + ) + + if val_loss < best_val_loss: + best_val_loss = val_loss + best_weights = model.get_weights() + patience_counter = es_patience + print(f" ↳ New best validation loss: {best_val_loss:.4f}") + else: + patience_counter -= 1 + if patience_counter == 0: + print("Early stopping triggered.") + break + + elapsed = time.perf_counter() - start_time + + # Restore best weights and evaluate on the test set. model.set_weights(best_weights) - test_loss, test_acc = evaluate(loader_test, model, loss_fn) - test_loss_r, test_acc_r = evaluate( - loader_test, model, loss_fn, retransform=True) + test_loader = _make_loader( + test_data, batch_size=min(batch_size, max(1, n_test)), shuffle=False) + test_loss, _ = evaluate(test_loader, model, loss_fn) + + test_loader_retransform = _make_loader( + test_data, batch_size=min(batch_size, max(1, n_test)), shuffle=False) + test_loss_orig, _ = evaluate( + test_loader_retransform, model, loss_fn, retransform=True) + + print(f"Test loss (log space): {test_loss:.4f}") + print(f"Test loss (original scale): {test_loss_orig:.4f}") + print(f"Training runtime: {elapsed:.2f}s ({elapsed / 60:.2f} min)") + + summary = TrainingSummary( + model_name=save_name, mean_train_loss=float( + np.min(epoch_train_losses)) + if epoch_train_losses else float("nan"), + mean_val_loss=float(np.min(epoch_val_losses)) + if epoch_val_losses else float("nan"), mean_test_loss=float(test_loss), + mean_test_loss_orig=float(test_loss_orig), + training_minutes=elapsed / 60, train_losses=[epoch_train_losses], + val_losses=[epoch_val_losses]) - print( - "Done. Test loss: {:.4f}. Test acc: {:.2f}".format( - test_loss, test_acc)) - test_scores.append(test_loss) - test_scores_r.append(test_loss_r) - train_losses.append(np.asarray(losses_history).min()) - val_losses.append(np.asarray(val_losses_history).min()) - losses_history_all.append(np.asarray(losses_history)) - val_losses_history_all.append(val_losses_history) - - elapsed = time.perf_counter() - start - - # print out stats - print(f"Best train losses: {train_losses} ") - print(f"Best validation losses: {val_losses}") - print(f"Test values: {test_scores}") - print("--------------------------------------------") - print(f"Train Score:{np.mean(train_losses)}") - print(f"Validation Score:{np.mean(val_losses)}") - print(f"Test Score (log): {np.mean(test_scores)}") - print(f"Test Score (orig.): {np.mean(test_scores_r)}") - - print(f"Time for training: {elapsed:.4f} seconds") - print(f"Time for training: {elapsed/60:.4f} minutes") - - ################################################################################ - # Save results - ################################################################################ if save_dir: - # save df - df.loc[len(df.index)] = [ # layer_name, number_of_layer, channels, - np.mean(train_losses), - np.mean(val_losses), - np.mean(test_scores), - np.mean(test_scores_r), - (elapsed / 60), - [losses_history_all], - [val_losses_history_all]] - print(df) - - # Ensure that save_name has the .pickle extension - if not save_name.endswith('.pickle'): - save_name += '.pickle' - - # Save best weights as pickle - file_path_w = os.path.join(save_dir, 'saved_weights') - # Ensure the directory exists - os.makedirs(file_path_w, exist_ok=True) - - # Construct the full file path by joining the directory with save_name - file_path_w = os.path.join(file_path_w, save_name) - - # Save the weights to the file - with open(file_path_w, 'wb') as f: - pickle.dump(best_weights, f) - - # Save evaluation CSV - file_path_df = os.path.join(save_dir, 'model_evaluations_paper') - os.makedirs(file_path_df, exist_ok=True) - file_path_df = os.path.join( - file_path_df, save_name.replace('.pickle', '.csv')) - df.to_csv(file_path_df) - - return { - "model": save_name, - "mean_train_loss": np.mean(train_losses), - "mean_val_loss": np.mean(val_losses), - "mean_test_loss": np.mean(test_scores), - "mean_test_loss_orig": np.mean(test_scores_r), - "training_time": elapsed/60, - "train_losses": losses_history_all, - "val_losses": val_losses_history_all - } + save_dir = Path(save_dir) + save_dir.mkdir(parents=True, exist_ok=True) + + metrics_df = pd.DataFrame(columns=[ + "train_loss", "val_loss", "test_loss", + "test_loss_orig", "training_time", + "loss_history", "val_loss_history"]) + metrics_df.loc[len(metrics_df.index)] = [ + summary.mean_train_loss, + summary.mean_val_loss, + summary.mean_test_loss, + summary.mean_test_loss_orig, + summary.training_minutes, + summary.train_losses, + summary.val_losses + ] + + weights_dir = save_dir / "saved_weights" + weights_dir.mkdir(parents=True, exist_ok=True) + + weights_filename = save_name if save_name.endswith( + ".pickle") else f"{save_name}.pickle" + weights_path = weights_dir / weights_filename + with weights_path.open("wb") as fp: + pickle.dump(best_weights, fp) + + results_dir = save_dir / "model_evaluations_paper" + results_dir.mkdir(parents=True, exist_ok=True) + results_path = results_dir / \ + weights_filename.replace(".pickle", ".csv") + metrics_df.to_csv(results_path, index=False) + print(f"Saved weights to {weights_path}") + print(f"Saved evaluation metrics to {results_path}") + + return asdict(summary) if __name__ == "__main__": @@ -337,15 +375,14 @@ def train_and_evaluate( optimizer = Adam(learning_rate=0.001) loss_fn = MeanAbsolutePercentageError() - path_memilio = os.path.abspath( - os.path.join( - os.path.dirname(__file__), '../../../../..')) + repo_root = Path(__file__).resolve().parents[4] + artifacts_root = repo_root / "artifacts" - # Generate the Dataset - path_cases = os.path.join(path_memilio, "saves", - "GNN_data_30days_3dampings_classic5.pickle") - path_mobility = os.path.join(path_memilio, "data", "Germany", "mobility") - data = create_dataset(path_cases, path_mobility) + dataset_path = artifacts_root / \ + "generated_datasets" / "GNN_data_30days_3dampings_classic5.pickle" + + mobility_dir = repo_root / "data" / "Germany" / "mobility" + data = load_gnn_dataset(dataset_path, mobility_dir) # Define the model architecture def transform_a(adjacency_matrix): @@ -366,9 +403,17 @@ def transform_a(adjacency_matrix): model = model_class() save_name = 'GNN_30days' # name for model + save_dir = artifacts_root / "model_results" + train_and_evaluate( - data, batch_size, epochs, model, loss_fn, optimizer, es_patience, - save_dir=os.path.join(path_memilio, "saves", "GNN_model_results"), + data, + batch_size, + epochs, + model, + loss_fn, + optimizer, + es_patience, + save_dir=save_dir, save_name=save_name) elapsed_hyper = time.perf_counter() - start_hyper From d39917d821138cd5a952b436ffc9699c8ccc04a5 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 14:05:26 +0100 Subject: [PATCH 27/41] . --- .../surrogatemodel/GNN/evaluate_and_train.py | 15 ++++--- .../test_surrogatemodel_GNN.py | 40 ++++++++++--------- 2 files changed, 32 insertions(+), 23 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index 6517030f37..75e8fce937 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -65,14 +65,19 @@ def __init__( self._node_labels = node_labels self._adjacency = adjacency.astype(np.float32) super().__init__() + # This must be set AFTER calling super().__init__() + self.a = self._adjacency def read(self): - """Create one Graph object per sample.""" + """Create one Graph object per sample. + + Note: For MixedLoader, individual Graph objects should not have an + adjacency matrix (a). The adjacency matrix is stored at the dataset level. + """ return [ spektral.data.Graph( x=feature.astype(np.float32), y=label.astype(np.float32), - a=self._adjacency, ) for feature, label in zip(self._node_features, self._node_labels) ] @@ -87,7 +92,7 @@ class TrainingSummary: mean_val_loss: float mean_test_loss: float mean_test_loss_orig: float - training_minutes: float + training_time: float train_losses: List[List[float]] val_losses: List[List[float]] @@ -326,7 +331,7 @@ def _make_loader(dataset, *, batch_size, shuffle=False): mean_val_loss=float(np.min(epoch_val_losses)) if epoch_val_losses else float("nan"), mean_test_loss=float(test_loss), mean_test_loss_orig=float(test_loss_orig), - training_minutes=elapsed / 60, train_losses=[epoch_train_losses], + training_time=elapsed / 60, train_losses=[epoch_train_losses], val_losses=[epoch_val_losses]) if save_dir: @@ -342,7 +347,7 @@ def _make_loader(dataset, *, batch_size, shuffle=False): summary.mean_val_loss, summary.mean_test_loss, summary.mean_test_loss_orig, - summary.training_minutes, + summary.training_time, summary.train_losses, summary.val_losses ] diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index e1bccad32e..92a5632083 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -56,11 +56,13 @@ def create_dummy_data( :param output_dim: Number of output dimensions per node. :return: A dictionary containing inputs, adjacency matrix, and labels. """ - X = np.random.rand(num_samples, 1, num_node_features, - num_nodes).astype(np.float32) + # Shape should be (num_samples, input_width, num_nodes, num_features) + X = np.random.rand(num_samples, 1, num_nodes, + num_node_features).astype(np.float32) A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) - y = np.random.rand(num_samples, 1, output_dim, - num_nodes).astype(np.float32) + # Shape should be (num_samples, label_width, num_nodes, label_features) + y = np.random.rand(num_samples, 1, num_nodes, + output_dim).astype(np.float32) return {"inputs": X, "adjacency": A, "labels": y} def setup_fake_filesystem(self, fs, path, data): @@ -227,11 +229,6 @@ def test_create_dataset(self): num_nodes = 5 num_node_features = 3 output_dim = 4 - X = np.random.rand(num_samples, 1, - num_node_features, num_nodes).astype(np.float32) - A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) - y = np.random.rand(num_samples, 1, - output_dim, num_nodes).astype(np.float32) data = self.create_dummy_data( num_samples, num_nodes, num_node_features, output_dim) # Save dummy data to the fake file system @@ -255,22 +252,29 @@ def test_create_dataset(self): def test_train_step(self): # Create a simple model for testing - model = tf.keras.Sequential([ - tf.keras.layers.Dense(10, activation='relu'), - tf.keras.layers.Dense(2) - ]) + model = get_model( + layer_type="GCNConv", num_layers=2, num_channels=16, + activation="relu", num_output=2) optimizer = tf.keras.optimizers.Adam() loss_fn = MeanAbsolutePercentageError() # Create dummy data - batch_size = 4 + num_samples = 10 num_nodes = 5 num_node_features = 3 output_dim = 2 - inputs = np.random.rand(batch_size, num_nodes, - num_node_features).astype(np.float32) - y = np.random.rand(batch_size, num_nodes, - output_dim).astype(np.float32) + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + # Build the model by calling it on a batch of data + loader = MixedLoader(dataset, batch_size=4, epochs=1) + inputs, y = next(loader) + model(inputs) # Perform a training step loss, acc = train_step(inputs, y, loss_fn, model, optimizer) From cf59960d661efdee780ea57dedede899243466f9 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 15:11:48 +0100 Subject: [PATCH 28/41] [ci skip] rm extrapolate gnn --- .../GNN/extrapolate_data_gnn.py | 195 ------------------ 1 file changed, 195 deletions(-) delete mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py deleted file mode 100644 index d33e7ad763..0000000000 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/extrapolate_data_gnn.py +++ /dev/null @@ -1,195 +0,0 @@ -############################################################################# -# Copyright (C) 2020-2025 MEmilio -# -# Authors: Henrik Zunker -# -# Contact: Martin J. Kuehn -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -############################################################################# - -import os -import numpy as np -import h5py -import pandas as pd -import memilio.simulation as mio -from memilio.simulation import (AgeGroup) - -from memilio.simulation.osecir import ( - InfectionState, Model, ModelGraph, set_nodes) -from memilio.surrogatemodel.GNN.GNN_utils import remove_confirmed_compartments -from memilio.surrogatemodel.GNN.data_generation import ( - set_contact_matrices, set_covid_parameters, Location) - - -def extrapolate_data(start_date, num_days, data_dir): - ''' - Extrapolate data using the C++ backend. - This function sets up an empty model and calls the C++ backend to proceed the real world - data. - :param start_date: The date to start the simulation from. - :param num_days: The number of days to simulate. - :param data_dir: The directory where the data is stored. - :return: None - ''' - # Set up the model - model = Model(6) - model.parameters.StartDay = start_date.day_in_year - set_covid_parameters(model) - set_contact_matrices(model, data_dir) - - graph = ModelGraph() - - # Set the parameters for the nodes - scaling_factor_infected = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - scaling_factor_icu = 1.0 - tnt_capacity_factor = 7.5 / 100000. - - # Path containing the population data - data_dir_Germany = os.path.join(data_dir, "Germany") - pydata_dir = os.path.join(data_dir_Germany, "pydata") - - path_population_data = os.path.join(pydata_dir, - "county_current_population.json") - # Set the nodes in the model - set_nodes( - model.parameters, - start_date, start_date + num_days, - pydata_dir, - path_population_data, True, graph, scaling_factor_infected, - scaling_factor_icu, tnt_capacity_factor, num_days, True) - - -def read_results_h5(path, group_key='Total'): - ''' Reads results from an HDF5 file and processes the data. - :param path: Path to the HDF5 file. - :param group_key: The key for the group to read from the HDF5 file. - :return: A dictionary with processed results. - ''' - with h5py.File(path, 'r') as f: - keys = list(f.keys()) - res = {} - for i, key in enumerate(keys): - group = f[key] - total = group[group_key][()] - # remove confirmed compartments - sum_inf_no_symp = np.sum(total[:, [2, 3]], axis=1) - sum_inf_symp = np.sum(total[:, [4, 5]], axis=1) - total[:, 2] = sum_inf_no_symp - total[:, 4] = sum_inf_symp - total = np.delete(total, [3, 5], axis=1) - res[key] = total - return res - - -def get_ground_truth_data(start_date, num_days, data_dir, create_new=False): - ''' - Generates ground truth data for the specified start date and number of days. - :param start_date: The date to start the simulation from. - :param num_days: The number of days to simulate. - :param data_dir: The directory where the data is stored. - :param create_new: If True, generates new data; otherwise, reads existing data. - :return: A numpy array containing the ground truth data. - ''' - # Check if the path exists and create new data if needed - path_rki_h5 = os.path.join(data_dir, "Germany", "pydata", "Results_rki.h5") - if not os.path.isfile(path_rki_h5) or create_new: - print("Generating real data from C++ backend...") - extrapolate_data(start_date, num_days, data_dir) - print("Data generation complete.") - - # age-stratified input. - num_age_groups = 6 - all_age_data_list = [] - group_keys = ['Group1', 'Group2', 'Group3', 'Group4', 'Group5', 'Group6'] - for age in range(num_age_groups): - age_data_dict = read_results_h5(path_rki_h5, group_keys[age]) - age_data_np = np.array(list(age_data_dict.values())) - all_age_data_list.append(age_data_np) - - all_age_data = np.stack(all_age_data_list, axis=-1) - num_nodes, timesteps, _, _ = all_age_data.shape - ground_truth_all_nodes_np = np.reshape( - all_age_data.transpose(0, 1, 3, 2), (num_nodes, timesteps, -1)) - - return ground_truth_all_nodes_np - - -def extrapolate_ground_truth_data(data_dir, num_days=360): - ''' - Extrapolates ground truth data for the specified number of days. - :param data_dir: The directory where the data is stored. - :param num_days: The number of days to simulate. - :return: None - ''' - cwd = os.getcwd() - data_dir = os.path.join(cwd, data_dir) - start_dates = [ - mio.Date(2020, 6, 1) - ] - - for start_date in start_dates: - get_ground_truth_data( - start_date, num_days, data_dir, create_new=True) - print(f"Ground truth data for {start_date} generated.") - - path_rki_h5 = data_dir + "/Germany/pydata/Results_rki.h5" - num_age_groups = 6 - all_age_data_list = [] - group_keys = ['Group1', 'Group2', 'Group3', 'Group4', 'Group5', 'Group6'] - for age in range(num_age_groups): - age_data_dict = read_results_h5(path_rki_h5, group_keys[age]) - age_data_np = np.array(list(age_data_dict.values())) - all_age_data_list.append(age_data_np) - all_age_data = np.stack(all_age_data_list, axis=-1) - - # Save the ground truth data - output_path = data_dir + "/Germany/pydata/ground_truth_all_nodes.pickle" - with open(output_path, 'wb') as f: - pd.to_pickle(all_age_data, f) - print(f"Ground truth data saved to {output_path}") - - -def generate_bounds(data_dir): - path = data_dir + "/Germany/pydata/ground_truth_all_nodes.pickle" - with open(path, 'rb') as f: - ground_truth_all_nodes_np = pd.read_pickle(f) - - # Calculate the bounds - lower_bound = np.min(ground_truth_all_nodes_np, axis=1) - upper_bound = np.max(ground_truth_all_nodes_np, axis=1) - path_upper_bound = data_dir + "/Germany/pydata/ground_truth_upper_bound.pickle" - path_lower_bound = data_dir + "/Germany/pydata/ground_truth_lower_bound.pickle" - with open(path_upper_bound, 'wb') as f: - pd.to_pickle(upper_bound, f) - with open(path_lower_bound, 'wb') as f: - pd.to_pickle(lower_bound, f) - - print(f"Upper bound saved to {path_upper_bound}") - print(f"Lower bound saved to {path_lower_bound}") - - -def main(): - cwd = os.getcwd() - num_days = 180 - data_dir = os.path.join(cwd, "data") - - extrapolate_ground_truth_data( - data_dir, - num_days - ) - generate_bounds(data_dir) - - -if __name__ == "__main__": - main() From d177e4c7ee631fefdf6bcea9c446cd08e101b7c2 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 15:12:14 +0100 Subject: [PATCH 29/41] [ci skip] rework gnn utils + grid_seach --- .../memilio/surrogatemodel/GNN/GNN_utils.py | 353 ++++++++++------ .../memilio/surrogatemodel/GNN/grid_search.py | 381 +++++++++++++----- 2 files changed, 506 insertions(+), 228 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py index a454b28359..e98d0ef99f 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -2,6 +2,7 @@ # Copyright (C) 2020-2025 MEmilio # # Authors: Agatha Schmidt, Henrik Zunker +# # Contact: Martin J. Kuehn # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -16,179 +17,279 @@ # See the License for the specific language governing permissions and # limitations under the License. ############################################################################# +""" +Utility functions for GNN-based surrogate models. + +This module provides helper functions for data preprocessing, transformation, +and graph construction used in Graph Neural Network surrogate models for +epidemiological simulations. +""" + +import os +from typing import List, Tuple import numpy as np import pandas as pd -import os from sklearn.preprocessing import FunctionTransformer -# memilio.epidata import transformMobilityData as tmd -# from memilio.epidata import getDataIntoPandasDataFrame as gd -# from memilio.simulation.osecir import (ModelGraph, set_edges) -# from memilio.epidata import modifyDataframeSeries as mdfs +from memilio.simulation.osecir import ModelGraph, set_edges +from memilio.epidata import getDataIntoPandasDataFrame as gd +from memilio.epidata import transformMobilityData as tmd +from memilio.epidata import modifyDataframeSeries as mdfs + + +# Default number of compartments in the ODE-SECIR model (without confirmed compartments) +DEFAULT_NUM_COMPARTMENTS = 8 def remove_confirmed_compartments(dataset_entries, num_groups): - """ The compartments which contain confirmed cases are not needed and are - therefore omitted by summarizing the confirmed compartment with the - original compartment. + """Removes confirmed compartments from simulation data by merging them with base compartments. + + The ODE-SECIR model includes separate "confirmed" compartments that track + diagnosed cases. For GNN training, these are merged back into their base + compartments (InfectedNoSymptoms + InfectedNoSymptomsConfirmed -> InfectedNoSymptoms, + InfectedSymptoms + InfectedSymptomsConfirmed -> InfectedSymptoms). - :param dataset_entries: Array that contains the compartmental data with - confirmed compartments. - :param num_groups: Number of age groups. - :returns: Array that contains the compartmental data without confirmed compartments. - """ + :param dataset_entries: Compartment data array containing confirmed compartments. + Shape: [num_timesteps, num_groups * num_compartments_with_confirmed] + :param num_groups: Number of age groups in the model. + :returns: Array with confirmed compartments merged into base compartments. + Shape: [num_timesteps, num_groups * num_compartments_without_confirmed] + + """ new_dataset_entries = [] - for i in dataset_entries: - dataset_entries_reshaped = i.reshape( - [num_groups, int(np.asarray(dataset_entries).shape[1]/num_groups)] + for timestep_data in dataset_entries: + # Reshape to separate age groups and compartments + data_reshaped = timestep_data.reshape( + [num_groups, int(np.asarray(dataset_entries).shape[1] / num_groups)] ) - sum_inf_no_symp = np.sum(dataset_entries_reshaped[:, [2, 3]], axis=1) - sum_inf_symp = np.sum(dataset_entries_reshaped[:, [4, 5]], axis=1) - dataset_entries_reshaped[:, 2] = sum_inf_no_symp - dataset_entries_reshaped[:, 4] = sum_inf_symp + + # Merge InfectedNoSymptoms (index 2) with InfectedNoSymptomsConfirmed (index 3) + sum_infected_no_symptoms = np.sum(data_reshaped[:, [2, 3]], axis=1) + + # Merge InfectedSymptoms (index 4) with InfectedSymptomsConfirmed (index 5) + sum_infected_symptoms = np.sum(data_reshaped[:, [4, 5]], axis=1) + + # Replace original compartments with merged values + data_reshaped[:, 2] = sum_infected_no_symptoms + data_reshaped[:, 4] = sum_infected_symptoms + + # Remove confirmed compartments (indices 3 and 5) and flatten new_dataset_entries.append( - np.delete(dataset_entries_reshaped, [3, 5], axis=1).flatten() + np.delete(data_reshaped, [3, 5], axis=1).flatten() ) + return new_dataset_entries -def getBaselineMatrix(): - """ loads the baselinematrix +def get_baseline_contact_matrix(data_dir): + """Loads and sums baseline contact matrices for all location types. + + Loads contact matrices for home, school, work, and other locations, then + returns their sum as the total baseline contact matrix. + + :param data_dir: Root directory containing contact matrix data. + :returns: Combined baseline contact matrix as numpy array. + :raises FileNotFoundError: If any contact matrix file is not found. + """ + contact_dir = os.path.join(data_dir, "Germany", "contacts") + + contact_files = [ + "baseline_home.txt", + "baseline_school_pf_eig.txt", + "baseline_work.txt", + "baseline_other.txt" + ] + + baseline_matrix = None + + for filename in contact_files: + filepath = os.path.join(contact_dir, filename) + + if not os.path.exists(filepath): + raise FileNotFoundError( + f"Contact matrix file not found: {filepath}") + + matrix = np.loadtxt(filepath) + + if baseline_matrix is None: + baseline_matrix = matrix + else: + baseline_matrix += matrix + + return baseline_matrix + + +def create_mobility_graph(mobility_dir, num_regions, county_ids, models): + """Creates a graph-ODE model with mobility connections between regions. + + Constructs a graph, where each node represents a region (county) with + its own ODE-SECIR model, and edges represent mobility flows between regions. + + :param mobility_dir: Directory containing mobility data files. + :param num_regions: Number of regions/counties to include in the graph. + :param county_ids: List of county IDs/keys for each region. + :param models: List of ODE-SECIR Model instances, one per county. + :returns: Configured graph with nodes and mobility edges. - baseline_contact_matrix0 = os.path.join( - "./data/Germany/contacts/baseline_home.txt") - baseline_contact_matrix1 = os.path.join( - "./data/Germany/contacts/baseline_school_pf_eig.txt") - baseline_contact_matrix2 = os.path.join( - "./data/Germany/contacts/baseline_work.txt") - baseline_contact_matrix3 = os.path.join( - "./data/Germany/contacts/baseline_other.txt") - - baseline = np.loadtxt(baseline_contact_matrix0) \ - + np.loadtxt(baseline_contact_matrix1) + \ - np.loadtxt(baseline_contact_matrix2) + \ - np.loadtxt(baseline_contact_matrix3) - - return baseline - - -def make_graph(directory, num_regions, countykey_list, models): - """ Generating graph with one node per region. - - Each node contains the county ID and a osecir model for the county. The edges - contain the mobility information between different counties/nodes. - - :param directory: Directory with mobility data. - :param num_regions: Number (int) of counties that should be added to the - graph-ODE model. Equals 400 for whole Germany. - :param countykey_list: List of keys/IDs for each county. - :models models: List of osecir Model with one model per county. - :returns: Graph-ODE model. - """ - from memilio.simulation.osecir import (ModelGraph, set_edges) + """ + # Initialize empty graph graph = ModelGraph() + + # Add one node per region with its model for i in range(num_regions): - graph.add_node(int(countykey_list[i]), models[i]) + graph.add_node(int(county_ids[i]), models[i]) - num_locations = 4 + # Number of contact locations (home, school, work, other) + num_contact_locations = 4 + + # Add mobility edges between nodes + parent_dir = os.path.abspath(os.path.join(mobility_dir, os.pardir)) + set_edges(parent_dir, graph, num_contact_locations) - set_edges(os.path.abspath(os.path.join(directory, os.pardir)), - graph, num_locations) return graph -def transform_mobility_directory(): - """ Transforms the mobility data by merging Eisenach and Wartburgkreis - """ - from memilio.epidata import transformMobilityData as tmd - from memilio.epidata import getDataIntoPandasDataFrame as gd - # get mobility data directory - arg_dict = gd.cli("commuter_official") +def transform_mobility_data(data_dir): + """Updates mobility data to merge Eisenach and Wartburgkreis counties. - directory = arg_dict['out_folder'].split('/pydata')[0] - directory = os.path.join(directory, 'Germany/mobility/') + :param data_dir: Root directory containing Germany mobility data. + :returns: Path to the updated mobility directory. + """ + # Get mobility data directory + mobility_dir = os.path.join(data_dir, 'Germany/mobility/') - # Merge Eisenach and Wartbugkreis in Input Data - tmd.updateMobility2022(directory, mobility_file='twitter_scaled_1252') + # Update mobility files by merging Eisenach and Wartburgkreis + tmd.updateMobility2022(mobility_dir, mobility_file='twitter_scaled_1252') tmd.updateMobility2022( - directory, mobility_file='commuter_mobility_2022') - return directory + mobility_dir, mobility_file='commuter_mobility_2022') + return mobility_dir -def get_population(): - """ - Loading the population data for the different counties and the different age groups. + +def load_population_data(data_dir, age_groups=None): + """Loads population data for counties stratified by age groups. + + Reads county-level population data and aggregates it into specified age groups. + Default age groups follow the standard 6-group structure used in the ODE-SECIR model. + + :param data_dir: Root directory containing population data (should contain Germany/pydata/ subdirectory). + :param age_groups: List of age group labels (default: ['0-4', '5-14', '15-34', '35-59', '60-79', '80-130']). + :returns: List of population data entries, where each entry is [county_id, pop_group1, pop_group2, ...]. + :raises FileNotFoundError: If population data file is not found. """ - from memilio.epidata import modifyDataframeSeries as mdfs - df_population = pd.read_json( - 'data/Germany/pydata/county_population.json') - age_groups = ['0-4', '5-14', '15-34', '35-59', '60-79', '80-130'] + if age_groups is None: + age_groups = ['0-4', '5-14', '15-34', '35-59', '60-79', '80-130'] + + # Load population data + population_file = os.path.join( + data_dir, 'Germany', 'pydata', 'county_population.json') + + if not os.path.exists(population_file): + raise FileNotFoundError( + f"Population data file not found: {population_file}") + + df_population = pd.read_json(population_file) + # Initialize DataFrame for age-grouped population df_population_agegroups = pd.DataFrame( - columns=[df_population.columns[0]] + age_groups) + columns=[df_population.columns[0]] + age_groups + ) + + # Process each region for region_id in df_population.iloc[:, 0]: - df_population_agegroups.loc[len(df_population_agegroups.index), :] = [int(region_id)] + list( - mdfs.fit_age_group_intervals(df_population[df_population.iloc[:, 0] == int(region_id)].iloc[:, 2:], age_groups)) + region_data = df_population[df_population.iloc[:, 0] + == int(region_id)] + age_grouped_pop = list( + mdfs.fit_age_group_intervals( + region_data.iloc[:, 2:], age_groups) + ) + + df_population_agegroups.loc[len(df_population_agegroups.index)] = [ + int(region_id)] + age_grouped_pop - population = df_population_agegroups.values.tolist() - return population + return df_population_agegroups.values.tolist() -def scale_data(data, transform=True, num_compartments=8): - """ Apply a logarithmic transformation on the data. +def scale_data( + data, transform=True, num_compartments=DEFAULT_NUM_COMPARTMENTS): + """Applies logarithmic transformation to simulation data for training. + + Transforms compartment data using log1p (log(1+x)) to stabilize training. + This transformation helps handle the wide range of population values and + improves gradient flow during neural network training. + + :param data: Dictionary containing 'inputs' and 'labels' keys with simulation data. + Inputs shape: [num_samples, time_steps, num_nodes, num_features] + Labels shape: [num_samples, time_steps, num_nodes, num_features] + :param transform: Whether to apply log transformation (True) or just reshape (False). + :param num_compartments: Number of compartments per age group in the model (default: 8). + :returns: Tuple of (scaled_inputs, scaled_labels), both with shape + [num_samples, num_features, time_steps, num_nodes] + :raises ValueError: If input data is not numeric or has unexpected structure. - :param data: dictionary, containing entries "inputs" and "labels" - :param transform: Boolean, if True apply the transformation, else return the data reshaped. - :param num_compartments: Number of compartments in the epidemiological model. - :returns scaled_inputs: Transformed input data - :returns scaled_labels: Transformed output data """ - if not np.issubdtype(np.asarray(data['inputs']).dtype, np.number): + # Validate input data types + inputs_array = np.asarray(data['inputs']) + labels_array = np.asarray(data['labels']) + + if not np.issubdtype(inputs_array.dtype, np.number): raise ValueError("Input data must be numeric.") - if not np.issubdtype(np.asarray(data['labels']).dtype, np.number): + if not np.issubdtype(labels_array.dtype, np.number): raise ValueError("Label data must be numeric.") - num_groups = int(np.asarray(data['inputs']).shape[2] / num_compartments) + + # Calculate number of age groups from data shape + num_groups = int(inputs_array.shape[2] / num_compartments) + + # Initialize transformer (log1p for numerical stability) transformer = FunctionTransformer(np.log1p, validate=True) - # Scale inputs - inputs = np.asarray( - data['inputs']).transpose(2, 0, 1, 3).reshape(num_groups * 8, -1) + # Process inputs + # Reshape: [samples, timesteps, nodes, features] -> [nodes, samples, timesteps, features] + # -> [nodes * compartments, samples * timesteps * age_groups] + inputs_reshaped = inputs_array.transpose(2, 0, 1, 3).reshape( + num_groups * num_compartments, -1 + ) + if transform: - scaled_inputs = transformer.transform(inputs) + inputs_transformed = transformer.transform(inputs_reshaped) else: - scaled_inputs = inputs - original_shape_input = np.asarray(data['inputs']).shape - - # Reverse the reshape - reshaped_back = scaled_inputs.reshape(original_shape_input[2], - original_shape_input[0], - original_shape_input[1], - original_shape_input[3]) - - # Reverse the transpose - original_inputs = reshaped_back.transpose(1, 2, 0, 3) - scaled_inputs = original_inputs.transpose(0, 3, 1, 2) - - # Scale labels - labels = np.asarray( - data['labels']).transpose(2, 0, 1, 3).reshape(num_groups * 8, -1) + inputs_transformed = inputs_reshaped + + original_shape_input = inputs_array.shape + + # Reverse reshape to separate dimensions + inputs_back = inputs_transformed.reshape( + original_shape_input[2], + original_shape_input[0], + original_shape_input[1], + original_shape_input[3] + ) + + # Reverse transpose and reorder to [samples, features, timesteps, nodes] + scaled_inputs = inputs_back.transpose(1, 2, 0, 3).transpose(0, 3, 1, 2) + + # Process labels with same procedure + labels_reshaped = labels_array.transpose(2, 0, 1, 3).reshape( + num_groups * num_compartments, -1 + ) + if transform: - scaled_labels = transformer.transform(labels) + labels_transformed = transformer.transform(labels_reshaped) else: - scaled_labels = labels - original_shape_labels = np.asarray(data['labels']).shape - - # Reverse the reshape - reshaped_back = scaled_labels.reshape(original_shape_labels[2], - original_shape_labels[0], - original_shape_labels[1], - original_shape_labels[3]) - - # Reverse the transpose - original_labels = reshaped_back.transpose(1, 2, 0, 3) - scaled_labels = original_labels.transpose(0, 3, 1, 2) + labels_transformed = labels_reshaped + + original_shape_labels = labels_array.shape + + labels_back = labels_transformed.reshape( + original_shape_labels[2], + original_shape_labels[0], + original_shape_labels[1], + original_shape_labels[3] + ) + + scaled_labels = labels_back.transpose(1, 2, 0, 3).transpose(0, 3, 1, 2) return scaled_inputs, scaled_labels diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py index f92ceccc23..194bb52909 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py @@ -17,23 +17,33 @@ # See the License for the specific language governing permissions and # limitations under the License. ############################################################################# +""" +Grid search module for GNN-based surrogate model hyperparameter optimization. + +This module provides functionality to perform systematic hyperparameter search +over GNN architectures, training configurations, and model parameters to identify +optimal configurations for epidemiological surrogate models. +""" import os -from typing import Optional +from pathlib import Path import pandas as pd import tensorflow as tf -from tensorflow.keras.losses import MeanAbsolutePercentageError as MeanAbsolutePercentageError -from tensorflow.keras.metrics import MeanAbsolutePercentageError as mean_absolute_percentage_error - -from tensorflow.keras.optimizers import Adam from spektral.data import MixedLoader -from memilio.surrogatemodel.GNN.network_architectures import get_model +from tensorflow.keras.losses import MeanAbsolutePercentageError +from tensorflow.keras.metrics import MeanAbsolutePercentageError as MeanAbsolutePercentageErrorMetric +from tensorflow.keras.optimizers import Adam + from memilio.surrogatemodel.GNN.evaluate_and_train import ( - train_and_evaluate, create_dataset) + create_dataset, + train_and_evaluate +) +from memilio.surrogatemodel.GNN.network_architectures import get_model -# Define the parameter grid for grid search -layer_types = [ + +# Default hyperparameter grid +DEFAULT_LAYER_TYPES = [ "ARMAConv", "GCSConv", "GATConv", @@ -41,107 +51,274 @@ "APPNPConv" ] -numbers_layers = [2, 3, 4, 5, 6, 7] -numbers_channels = [2, 3, 4, 5, 6, 7] -activation_functions = ["elu", "relu", "tanh", "sigmoid"] -model_parameters = [ - (layer, num_layers, num_channels, activation) - for layer in layer_types - for num_layers in numbers_layers - for num_channels in numbers_channels for activation in - activation_functions] +DEFAULT_NUM_LAYERS = [2, 3, 4, 5, 6, 7] +DEFAULT_NUM_CHANNELS = [2, 3, 4, 5, 6, 7] +DEFAULT_ACTIVATION_FUNCTIONS = ["elu", "relu", "tanh", "sigmoid"] + +# Default training configuration +DEFAULT_BATCH_SIZE = 32 +DEFAULT_MAX_EPOCHS = 100 +DEFAULT_ES_PATIENCE = 30 + -# Fix training parameters -batch_size = 32 -loss_function = MeanAbsolutePercentageError() -optimizer = Adam() -es_patience = 30 -max_epochs = 100 +def generate_parameter_grid( + layer_types=None, + num_layers_options=None, + num_channels_options=None, + activation_functions=None): + """Generates a grid of model parameter combinations for hyperparameter search. + + :param layer_types: List of GNN layer types to test (default: DEFAULT_LAYER_TYPES). + :param num_layers_options: List of layer counts to test (default: DEFAULT_NUM_LAYERS). + :param num_channels_options: List of channel counts to test (default: DEFAULT_NUM_CHANNELS). + :param activation_functions: List of activation functions to test (default: DEFAULT_ACTIVATION_FUNCTIONS). + :returns: List of tuples, each containing (layer_type, num_layers, num_channels, activation). + + """ + if layer_types is None: + layer_types = DEFAULT_LAYER_TYPES + if num_layers_options is None: + num_layers_options = DEFAULT_NUM_LAYERS + if num_channels_options is None: + num_channels_options = DEFAULT_NUM_CHANNELS + if activation_functions is None: + activation_functions = DEFAULT_ACTIVATION_FUNCTIONS -training_parameters = [batch_size, loss_function, - optimizer, es_patience, max_epochs] + parameter_grid = [ + (layer_type, num_layers, num_channels, activation) + for layer_type in layer_types + for num_layers in num_layers_options + for num_channels in num_channels_options + for activation in activation_functions + ] + + return parameter_grid def perform_grid_search( - model_parameters, training_parameters, data, - save_dir: Optional[str] = None): - """Perform grid search over the specified model parameters. - Trains and evaluates models with different configurations and stores the results in a CSV file. - Results are saved in 'grid_search_results.csv' in the 'saves' directory. - :param model_parameters: List of tuples containing model parameters (layer_type, num_layers, num_channels, activation). - :param training_parameters: List containing training parameters (batch_size, loss_function, optimizer, es_patience, max_epochs). - :param data: The dataset to be used for training and evaluation. - :return: None + data, + parameter_grid, + save_dir, + batch_size=DEFAULT_BATCH_SIZE, + max_epochs=DEFAULT_MAX_EPOCHS, + es_patience=DEFAULT_ES_PATIENCE, + learning_rate=0.001): + """Performs systematic grid search over GNN hyperparameters. + + Trains and evaluates models for each parameter combination in the grid, + tracking performance metrics and saving results. + Results are stored in a CSV file for later analysis. + + :param data: Spektral dataset containing training, validation, and test samples. + :param parameter_grid: List of tuples with (layer_type, num_layers, num_channels, activation) combinations. + :param save_dir: Directory to save results. + :param batch_size: Batch size for training (default: 32). + :param max_epochs: Maximum number of training epochs per configuration (default: 100). + :param es_patience: Early stopping in epochs (default: 30). + :param learning_rate: Learning rate for Adam optimizer (default: 0.001). + :returns: DataFrame containing all grid search results. + :raises ValueError: If data is empty or parameter_grid is invalid. + """ - # Setting hyprparameter + if not data or len(data) == 0: + raise ValueError("Dataset must contain at least one sample.") + + if not parameter_grid or len(parameter_grid) == 0: + raise ValueError( + "Parameter grid must contain at least one configuration.") + + # Determine output dimension from data output_dim = data[0].y.shape[-1] - batch_size, loss_function, optimizer, es_patience, max_epochs = training_parameters - # Create a DataFrame to store the results - df_results = pd.DataFrame( - columns=['model', 'optimizer', - 'number_of_hidden_layers', - 'number_of_channels', 'activation', - 'mean_train_loss', - 'mean_validation_loss', 'training_time', - 'train_losses', 'val_losses']) - - for param in model_parameters: - layer_type, num_layers, num_channels, activation = param - print( - f"Training model with {layer_type}, {num_layers} layers, {num_channels} channels, activation: {activation}") - - # Create a model instance - model = get_model( - layer_type=layer_type, - num_layers=num_layers, - num_channels=num_channels, - activation=activation, - num_output=output_dim - ) - optimizer = Adam() - - loader = MixedLoader(data, batch_size=batch_size, epochs=1) - inputs, _ = loader.__next__() - model(inputs) # Build the model by calling it on a batch of data - - # Initialize optimizer variables - optimizer.build(model.trainable_variables) - model.compile( - optimizer=optimizer, - loss=MeanAbsolutePercentageError(), - metrics=[mean_absolute_percentage_error()] - ) - - results = train_and_evaluate( - data, batch_size, epochs=max_epochs, model=model, - loss_fn=loss_function, optimizer=optimizer, - es_patience=es_patience, save_name="", save_dir=save_dir) - - df_results.loc[len(df_results.index)] = [ - layer_type, optimizer.__class__.__name__, num_layers, num_channels, activation, - results["mean_train_loss"], - results["mean_val_loss"], - results["training_time"], - results["train_losses"], - results["val_losses"] + + # Initialize loss function and optimizer + loss_function = MeanAbsolutePercentageError() + + # Initialize results DataFrame + results_df = pd.DataFrame( + columns=[ + 'model', + 'optimizer', + 'number_of_hidden_layers', + 'number_of_channels', + 'activation', + 'mean_train_loss', + 'mean_validation_loss', + 'mean_test_loss', + 'mean_test_loss_orig', + 'training_time', + 'train_losses', + 'val_losses' ] - # Save intermediate results to avoid data loss - # If save_dir is provided, use it. Otherwise store under a local 'saves' folder next to this module - base_dir = save_dir if ( - save_dir and len(save_dir) > 0) else os.path.realpath( - os.path.dirname(__file__)) - saves_dir = os.path.join(base_dir, 'saves') - os.makedirs(saves_dir, exist_ok=True) - df_results.to_csv(os.path.join( - saves_dir, 'grid_search_results.csv'), index=False) - print(f"Saved intermediate results to {saves_dir}") - # Clear session to free memory after each iteration - tf.keras.backend.clear_session() + ) + + save_dir.mkdir(parents=True, exist_ok=True) + results_file = save_dir / 'grid_search_results.csv' + + print(f"\n{'=' * 70}") + print("GNN Grid Search - Hyperparameter Optimization") + print(f"{'=' * 70}") + print(f"Total configurations to test: {len(parameter_grid)}") + print(f"Results will be saved to: {results_file}") + print(f"{'=' * 70}\n") + + # Iterate through all parameter combinations + for idx, (layer_type, num_layers, num_channels, activation) in enumerate( + parameter_grid, 1): + print(f"\n[{idx}/{len(parameter_grid)}] Training configuration:") + print(f" Layer type: {layer_type}") + print(f" Number of layers: {num_layers}") + print(f" Number of channels: {num_channels}") + print(f" Activation function: {activation}") + + try: + # Create model instance + model = get_model( + layer_type=layer_type, + num_layers=num_layers, + num_channels=num_channels, + activation=activation, + num_output=output_dim + ) + + # Initialize optimizer for this configuration + optimizer = Adam(learning_rate=learning_rate) + + # Build model by passing a sample batch through it + build_loader = MixedLoader( + data, batch_size=batch_size, epochs=1, shuffle=False) + build_inputs, _ = next(build_loader) + model(build_inputs) + + # Initialize optimizer variables + optimizer.build(model.trainable_variables) + + model.compile( + optimizer=optimizer, + loss=MeanAbsolutePercentageError(), + metrics=[MeanAbsolutePercentageErrorMetric()] + ) + + # Train and evaluate model + results = train_and_evaluate( + data=data, + batch_size=batch_size, + epochs=max_epochs, + model=model, + loss_fn=loss_function, + optimizer=optimizer, + es_patience=es_patience, + save_name="", # Dont save individual models during grid search + save_dir=None # Dont save individual results + ) + + # Store results + results_df.loc[len(results_df)] = [ + layer_type, + optimizer.__class__.__name__, + num_layers, + num_channels, + activation, + results["mean_train_loss"], + results["mean_val_loss"], + results["mean_test_loss"], + results["mean_test_loss_orig"], + results["training_time"], + results["train_losses"], + results["val_losses"] + ] + + # Save intermediate results after each configuration + results_df.to_csv(results_file, index=False) + print( + f" ✓ Configuration complete. Results saved to {results_file}") + + except Exception as e: + print(f" ✗ Error training configuration: {e}") + # Continue with next configuration rather than failing entire search + + finally: + # Clear TensorFlow session to free memory + tf.keras.backend.clear_session() + + print(f"\n{'=' * 70}") + print("Grid Search Complete!") + print(f"{'=' * 70}") + print(f"Total configurations tested: {len(results_df)}") + print(f"Results saved to: {results_file}") + + # Print best configuration + if len(results_df) > 0: + best_idx = results_df['mean_val_loss'].idxmin() + best_config = results_df.loc[best_idx] + print(f"\nBest Configuration:") + print(f" Model: {best_config['model']}") + print(f" Layers: {best_config['number_of_hidden_layers']}") + print(f" Channels: {best_config['number_of_channels']}") + print(f" Activation: {best_config['activation']}") + print(f" Validation Loss: {best_config['mean_validation_loss']:.4f}") + print(f" Test Loss: {best_config['mean_test_loss']:.4f}") + + print(f"{'=' * 70}\n") + + return results_df + + +def main(): + """Main function demonstrating grid search usage. + """ + # Dataset path + dataset_path = Path.cwd() / "generated_datasets" / \ + "GNN_data_30days_3dampings_classic5.pickle" + mobility_dir = Path.cwd() / "data" / "Germany" / "mobility" + + # Output directory for results + output_dir = Path.cwd() / "grid_search_results" + + print("=" * 70) + print("GNN Grid Search - Configuration") + print("=" * 70) + print(f"Dataset: {dataset_path}") + print(f"Mobility data: {mobility_dir}") + print(f"Output directory: {output_dir}") + print("=" * 70) + + # Verify files exist + if not dataset_path.exists(): + raise FileNotFoundError(f"Dataset not found: {dataset_path}") + if not mobility_dir.exists(): + raise FileNotFoundError( + f"Mobility directory not found: {mobility_dir}") + + # Load dataset + print("\nLoading dataset...") + data = create_dataset(str(dataset_path), str(mobility_dir)) + print(f"Dataset loaded: {len(data)} samples") + + # Generate parameter grid + # For demonstration, use a smaller grid. Remove restrictions for full search. + parameter_grid = generate_parameter_grid( + layer_types=["ARMAConv", "GCNConv"], + num_layers_options=[3, 5], + num_channels_options=[4, 6], + activation_functions=["elu", "relu"] + ) + + print( + f"Generated parameter grid with {len(parameter_grid)} configurations") + + # Perform grid search + results = perform_grid_search( + data=data, + parameter_grid=parameter_grid, + save_dir=str(output_dir), + batch_size=32, + max_epochs=100, + es_patience=30, + learning_rate=0.001 + ) + + print(f"\nGrid search complete. Results shape: {results.shape}") if __name__ == "__main__": - # Generate the Dataset - path_cases = "/localdata1/hege_mn/memilio/saves/GNN_data_30days_3dampings_classic5.pickle" - path_mobility = '/localdata1/hege_mn/memilio/data/Germany/mobility' - data = create_dataset(path_cases, path_mobility) - perform_grid_search(model_parameters, data) + main() From 223b3d53ea2be78d9e0c8e4331e04de14bd55ec7 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 16:00:55 +0100 Subject: [PATCH 30/41] . --- .../memilio/surrogatemodel/GNN/grid_search.py | 10 +- .../GNN/network_architectures.py | 434 +++++++++++++----- .../test_surrogatemodel_GNN.py | 37 +- 3 files changed, 354 insertions(+), 127 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py index 194bb52909..b6f78db540 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/grid_search.py @@ -25,7 +25,6 @@ optimal configurations for epidemiological surrogate models. """ -import os from pathlib import Path import pandas as pd @@ -127,6 +126,9 @@ def perform_grid_search( raise ValueError( "Parameter grid must contain at least one configuration.") + # Convert save_dir to Path if it's a string + save_dir = Path(save_dir) / "saves" + # Determine output dimension from data output_dim = data[0].y.shape[-1] @@ -230,10 +232,10 @@ def perform_grid_search( # Save intermediate results after each configuration results_df.to_csv(results_file, index=False) print( - f" ✓ Configuration complete. Results saved to {results_file}") + f"Configuration complete. Results saved to {results_file}") except Exception as e: - print(f" ✗ Error training configuration: {e}") + print(f"Error training configuration: {e}") # Continue with next configuration rather than failing entire search finally: @@ -248,7 +250,7 @@ def perform_grid_search( # Print best configuration if len(results_df) > 0: - best_idx = results_df['mean_val_loss'].idxmin() + best_idx = results_df['mean_validation_loss'].idxmin() best_config = results_df.loc[best_idx] print(f"\nBest Configuration:") print(f" Model: {best_config['model']}") diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py index 8a8b384261..55321d5b00 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/network_architectures.py @@ -17,183 +17,411 @@ # See the License for the specific language governing permissions and # limitations under the License. ############################################################################# +""" +Network architecture module for GNN-based surrogate models. + +This file provides functionality to generate Graph Neural Network +architectures with various layer types, configurations, and preprocessing +transformations. +""" + import numpy as np import tensorflow as tf -import spektral -import spektral.layers as spektral_layers -import spektral.utils.convolution as spektral_convolution +from spektral import layers as spektral_layers +from spektral.utils import convolution as spektral_convolution + + +# Supported GNN layer types +SUPPORTED_LAYER_TYPES = [ + "ARMAConv", + "GCSConv", + "GATConv", + "GCNConv", + "APPNPConv" +] + + +def _rescale_laplacian(laplacian): + """Rescales a Laplacian matrix while remaining compatible with newer SciPy.""" + laplacian = laplacian.toarray() if hasattr( + laplacian, "toarray") else np.asarray(laplacian) + try: + return spektral_convolution.rescale_laplacian(laplacian) + except TypeError as err: + if "eigvals" not in str(err): + raise + + lmax = np.linalg.eigvalsh(laplacian)[-1] + if lmax <= 0: + lmax = 2.0 + identity = np.eye(laplacian.shape[0], dtype=laplacian.dtype) + return (2.0 / lmax) * laplacian - identity + + +def _apply_graph_transform(adjacency, single_transform): + """Applies a transform function to adjacency tensors with rank 2 or 3.""" + adj_array = adjacency.numpy() + + if adj_array.ndim == 2: + transformed = np.asarray(single_transform(adj_array), dtype=np.float32) + return tf.convert_to_tensor(transformed, dtype=tf.float32) + if adj_array.ndim == 3: + transformed = [ + np.asarray(single_transform(graph_adj), dtype=np.float32) + for graph_adj in adj_array + ] + return tf.convert_to_tensor(np.stack(transformed), dtype=tf.float32) + + raise ValueError( + f"Adjacency tensor must be rank 2 or 3, got rank {adj_array.ndim}." + ) -# Function to generate a model class dynamically def generate_model_class( - name, layer_types, num_repeat, num_output, transform=None): - '''Generates a custom Keras model class with specified layers. + name, + layer_types, + num_repeat, + num_output, + transform=None): + """Dynamically generates a custom Keras GNN model class with specified layer configuration. + + This function creates a new Keras Model class with a configurable sequence of layers. + Each layer type can be repeated multiple times, allowing for flexible architecture design. + The generated model supports both graph-based layers (Spektral) and standard layers. - :param name: Name of the generated class. - :param layer_types: List of layer types (classes or callables) to include in the model. - :param num_repeat: List of integers specifying how many times to repeat each layer type. - :param num_output: Number of output units in the final layer. - :param transform: Optional function to transform the adjacency matrix before passing it to layers. - :return: A dynamically created Keras model class.''' + :param name: Name for the generated model class. + :param layer_types: List of layer constructors. + Each element should be a callable that instantiates a layer. + :param num_repeat: List of integers specifying repetition count for each layer type. + Must have same length as layer_types. + :param num_output: Number of output units in the final dense layer. + :param transform: Optional function to preprocess the adjacency matrix before passing + it to graph layers. Should accept a tensor and return a tensor. + :returns: Created Keras Model class. + :raises ValueError: If layer_types and num_repeat have different lengths, or if + any num_repeat value is less than 1. + + """ def __init__(self): + """Initializes the model with the specified layer sequence.""" if len(layer_types) != len(num_repeat): raise ValueError( - "layer_types and num_repeat must have the same length.") + f"layer_types and num_repeat must have the same length. " + f"Got {len(layer_types)} and {len(num_repeat)}." + ) if any(n < 1 for n in num_repeat): - raise ValueError("All values in num_repeat must be at least 1.") + raise ValueError( + "All values in num_repeat must be at least 1." + ) super(type(self), self).__init__() + + # Build sequence of hidden layers self.layer_seq = [] - for i, layer_type in enumerate(layer_types): - for _ in range(num_repeat[i]): + for layer_idx, layer_type in enumerate(layer_types): + for _ in range(num_repeat[layer_idx]): + # Instantiate layer from callable layer = layer_type() if callable(layer_type) else layer_type self.layer_seq.append(layer) - # Final output layer + + # Final output layer with ReLU activation self.output_layer = tf.keras.layers.Dense( - num_output, activation="relu") + num_output, activation="relu" + ) def call(self, inputs, mask=None): - # Input consists of a tuple (x, a) where x is the node features and a is the adjacency matrix + """Forward pass through the model. + + :param inputs: Tuple of (node_features, adjacency_matrix) where: + - node_features: [batch_size, num_nodes, num_features] + - adjacency_matrix: [num_nodes, num_nodes] or [batch_size, num_nodes, num_nodes] + :param mask: Optional node mask from data loader. Can be None or a list [node_mask, None]. + :returns: Model output tensor with shape [batch_size, num_nodes, num_output]. + + """ + # Unpack inputs x, a = inputs - # Extract potential node mask coming from the data loader and ensure broadcastable shape + # Extract and prepare node mask for masking operations node_mask = None if mask is not None: - # When inputs is a list [x, a], Keras/Spektral typically forwards a list of masks [node_mask, None] - node_mask = mask[0] if isinstance(mask, (tuple, list)) else mask + # Spektral typically provides masks as [node_mask, None] for [x, a] + node_mask = mask[0] if isinstance( + mask, (tuple, list)) else mask + if node_mask is not None: - # Expected by Spektral: shape [B, N, 1] to broadcast with [B, N, C] + # Ensure mask has shape [batch_size, num_nodes, 1] if tf.rank(node_mask) == 2: node_mask = tf.expand_dims(node_mask, axis=-1) - # If no node_mask was provided by the pipeline, create a broadcastable ones-mask + + # Create default mask if none provided if node_mask is None: - # x has shape [B, N, F] + # x has shape [batch_size, num_nodes, features] x_shape = tf.shape(x) node_mask = tf.ones([x_shape[0], x_shape[1], 1], dtype=tf.float32) - # Apply transformation on adjacency matrix if provided + # Apply adjacency matrix transformation if provided if not tf.is_symbolic_tensor(a): if transform is not None: a = transform(a) - # Pass through the layers + # Forward pass through layer sequence for layer in self.layer_seq: + # Check if layer is a Spektral graph layer if type(layer).__module__.startswith("spektral.layers"): - # Pass both `x` and `a` to the layer + # Graph layers need both features and adjacency matrix x = layer([x, a], mask=[node_mask, None]) else: - x = layer(x) # Pass only `x` to the layer + # Standard layers only need features + x = layer(x) + # Apply final output layer output = self.output_layer(x) return output - # Define the methods + # Create class dictionary with methods class_dict = { '__init__': __init__, 'call': call } + return type(name, (tf.keras.Model,), class_dict) -def get_model(layer_type, num_layers, num_channels, activation, num_output=1): - """Generates a GNN model based on the specified parameters. - :param layer_type: Name of GNN layer to use (possible 'ARMAConv', 'GCSConv', 'GATConv', - 'GCNConv', 'APPNPConv'), provided as a string. - :param num_layers: Number of hidden layers in the model. - :param num_channels: Number of channels (units) in each hidden layer. - :param activation: Activation function to use in the hidden layers (e.g., 'relu', 'elu', 'tanh', 'sigmoid'). - :param num_output: Number of output units in the final layer. - :return: A Keras model instance with the specified architecture. - """ - if layer_type not in [ - "ARMAConv", "GCSConv", "GATConv", "GCNConv", "APPNPConv"]: - raise ValueError( - f"Unsupported layer_type: {layer_type}. Supported types are 'ARMAConv', 'GCSConv', 'GATConv', 'GCNConv', 'APPNPConv'.") - if num_layers < 1: - raise ValueError("num_layers must be at least 1.") - if num_channels < 1: - raise ValueError("num_channels must be at least 1.") - if not isinstance(activation, str): - raise ValueError( - "activation must be a string representing the activation function.") - if num_output < 1: - raise ValueError("num_output must be at least 1.") +def _get_layer_config(layer_type): + """Returns layer class and transformation function for a given layer type. + + Each GNN layer type requires specific preprocessing of the adjacency matrix. + This function encapsulates the layer-specific configuration. - # Define the layer based on the specified type + :param layer_type: String identifier for the GNN layer type. + :returns: Tuple of (layer_class, transform_function). + :raises ValueError: If layer_type is not supported. + + """ if layer_type == "ARMAConv": - layer_name = spektral_layers.ARMAConv + layer_class = spektral_layers.ARMAConv + + def transform(adjacency): + """Applies rescaled Laplacian transformation for ARMA convolution.""" + + def single_transform(adj_array): + laplacian = spektral_convolution.normalized_laplacian( + adj_array) + return _rescale_laplacian(laplacian) + + return _apply_graph_transform(adjacency, single_transform) + + return layer_class, transform - def transform(a): - a = np.array(a) - a = spektral_convolution.rescale_laplacian( - spektral_convolution.normalized_laplacian(a)) - return tf.convert_to_tensor(a, dtype=tf.float32) elif layer_type == "GCSConv": - layer_name = spektral_layers.GCSConv + layer_class = spektral_layers.GCSConv + + def transform(adjacency): + """Applies normalized adjacency for GCS convolution.""" + return _apply_graph_transform( + adjacency, + spektral_convolution.normalized_adjacency + ) + + return layer_class, transform - def transform(a): - a = a.numpy() - a = spektral_convolution.normalized_adjacency(a) - return tf.convert_to_tensor(a, dtype=tf.float32) elif layer_type == "GATConv": - layer_name = spektral_layers.GATConv + layer_class = spektral_layers.GATConv + + def transform(adjacency): + """Applies normalized adjacency for GAT convolution.""" + return _apply_graph_transform( + adjacency, + spektral_convolution.normalized_adjacency + ) + + return layer_class, transform - def transform(a): - a = a.numpy() - a = spektral_convolution.normalized_adjacency(a) - return tf.convert_to_tensor(a, dtype=tf.float32) elif layer_type == "GCNConv": - layer_name = spektral_layers.GCNConv + layer_class = spektral_layers.GCNConv + + def transform(adjacency): + """Applies GCN filter for GCN convolution.""" + return _apply_graph_transform( + adjacency, + spektral_convolution.gcn_filter + ) + + return layer_class, transform - def transform(a): - a = a.numpy() - a = spektral_convolution.gcn_filter(a) - return tf.convert_to_tensor(a, dtype=tf.float32) elif layer_type == "APPNPConv": - layer_name = spektral_layers.APPNPConv + layer_class = spektral_layers.APPNPConv - def transform(a): - a = a.numpy() - a = spektral_convolution.gcn_filter(a) - return tf.convert_to_tensor(a, dtype=tf.float32) - # Generate the input for generate_model_class - layer_types = [ - lambda: - layer_name( - num_channels, activation=activation, - kernel_initializer=tf.keras.initializers.GlorotUniform())] + def transform(adjacency): + """Applies GCN filter for APPNP convolution.""" + return _apply_graph_transform( + adjacency, + spektral_convolution.gcn_filter + ) + + return layer_class, transform + + else: + raise ValueError( + f"Unsupported layer_type: '{layer_type}'. " + f"Supported types are: {', '.join(SUPPORTED_LAYER_TYPES)}" + ) + + +def get_model( + layer_type, + num_layers, + num_channels, + activation, + num_output=1): + """Generates a GNN model instance with specified architecture. + + Creates a Graph Neural Network model with a layer type repeated + multiple times. The model includes appropriate preprocessing for the adjacency + matrix based on the layer type. + + :param layer_type: Type of GNN layer to use. Must be one of: 'ARMAConv', 'GCSConv', + 'GATConv', 'GCNConv', 'APPNPConv'. + :param num_layers: Number of hidden GNN layers to stack. + :param num_channels: Number of channels (units/features) in each hidden layer. + :param activation: Activation function for hidden layers (e.g., 'relu', 'elu', 'tanh', 'sigmoid'). + :param num_output: Number of output units in the final dense layer (default: 1). + :returns: Instantiated Keras Model ready for training. + :raises ValueError: If parameters are invalid or layer_type is not supported. + + """ + # Validate inputs + if layer_type not in SUPPORTED_LAYER_TYPES: + raise ValueError( + f"Unsupported layer_type: '{layer_type}'. " + f"Supported types are: {', '.join(SUPPORTED_LAYER_TYPES)}" + ) + + if num_layers < 1: + raise ValueError( + f"num_layers must be at least 1, got {num_layers}." + ) + + if num_channels < 1: + raise ValueError( + f"num_channels must be at least 1, got {num_channels}." + ) + + if not isinstance(activation, str): + raise ValueError( + f"activation must be a string, got {type(activation).__name__}." + ) + + if num_output < 1: + raise ValueError( + f"num_output must be at least 1, got {num_output}." + ) + + # Get layer configuration for the specified type + layer_class, transform_fn = _get_layer_config(layer_type) + + # Create layer constructor with specified parameters + def create_layer(): + return layer_class( + num_channels, + activation=activation, + kernel_initializer=tf.keras.initializers.GlorotUniform() + ) + + # Configure model structure + layer_types = [create_layer] num_repeat = [num_layers] + + # Generate model class model_class = generate_model_class( - "CustomModel", layer_types, num_repeat, num_output, - transform=transform) + name="CustomGNNModel", + layer_types=layer_types, + num_repeat=num_repeat, + num_output=num_output, + transform=transform_fn + ) + return model_class() -if __name__ == "__main__": +def main(): + """Main function demonstrating network architecture usage. + """ + print("=" * 70) + print("GNN Network Architecture - Examples") + print("=" * 70) + + # Example 1: Custom model with mixed layer types + print("\nExample 1: Custom model with mixed Dense layers") layer_types = [ - # Dense layer (only uses x) lambda: tf.keras.layers.Dense(10, activation="relu"), - # Dense layer (only uses x) lambda: tf.keras.layers.Dense(20, activation="relu"), lambda: tf.keras.layers.Dense(30, activation="relu") ] - num_repeat = [2, 3, 1] + num_repeat = [2, 3, 1] # 2 layers of 10 units, 3 of 20, 1 of 30 + + custom_model_class = generate_model_class( + name="CustomMixedModel", + layer_types=layer_types, + num_repeat=num_repeat, + num_output=2 + ) + model1 = custom_model_class() + + # Example inputs for dense layers (no graph structure) + batch_size = 8 + num_nodes = 20 + num_features = 5 - testclass = generate_model_class( - "testclass", layer_types, num_repeat, num_output=2) - model = testclass() + # Node features: [batch_size, num_nodes, features] + x1 = tf.random.normal([batch_size, num_nodes, num_features]) + # Adjacency matrix per sample (not used by dense layers but required for input signature) + a1_single = tf.random.normal([num_nodes, num_nodes]) + a1 = tf.tile(tf.expand_dims(a1_single, axis=0), [batch_size, 1, 1]) + # Labels: [batch_size, num_nodes, outputs] + labels1 = tf.random.normal([batch_size, num_nodes, 2]) - # Example inputs - x = tf.random.normal([10, 20]) # Node features - a = tf.random.normal([10, 10]) # Adjacency matrix - labels = tf.random.normal([10, 2]) # Example labels + print("Compiling and training custom model...") + # Build the model by calling it once + _ = model1([x1, a1]) + model1.compile(optimizer="adam", loss="mse") + model1.fit([x1, a1], labels1, epochs=5, verbose=0) + print("Custom model trained successfully") - model.compile(optimizer="adam", loss="mse") - model.fit([x, a], labels, epochs=5) + # Example 2: Pre-configured GNN model + print("\nExample 2: Pre-configured ARMA GNN model") + model2 = get_model( + layer_type="ARMAConv", + num_layers=3, + num_channels=16, + activation="relu", + num_output=2 + ) - model2 = get_model("ARMAConv", 3, 16, "relu", num_output=2) + # Example inputs for GNN (proper graph structure) + # Node features: [batch_size, num_nodes, features] + x2 = tf.random.normal([batch_size, num_nodes, num_features]) + # Adjacency matrix per sample + a2_single = tf.eye( + num_nodes) + 0.1 * tf.random.normal([num_nodes, num_nodes]) + a2 = tf.tile(tf.expand_dims(a2_single, axis=0), [batch_size, 1, 1]) + # Labels: [batch_size, num_nodes, outputs] + labels2 = tf.random.normal([batch_size, num_nodes, 2]) + print("Compiling and training ARMA model...") + # Build the model by calling it once + _ = model2([x2, a2]) model2.compile(optimizer="adam", loss="mse") - model2.fit([x, a], labels, epochs=5) + model2.fit([x2, a2], labels2, epochs=5, verbose=0) + print("ARMA model trained successfully") + + print("\n" + "=" * 70) + print("Examples completed successfully!") + print("=" * 70 + "\n") + + +if __name__ == "__main__": + main() diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py index 92a5632083..58419294fb 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel_test/test_surrogatemodel_GNN.py @@ -140,7 +140,8 @@ def test_generate_model_class(self): model = ModelClass() self.assertEqual( str(error.exception), - "layer_types and num_repeat must have the same length.") + "layer_types and num_repeat must have the same length. " + "Got 1 and 2.") # Test with multiple layer types layer_types = [ @@ -186,8 +187,8 @@ def test_get_model(self): num_channels, activation, num_output) self.assertEqual( str(error.exception), - "Unsupported layer_type: MonvConv. " - "Supported types are 'ARMAConv', 'GCSConv', 'GATConv', 'GCNConv', 'APPNPConv'.") + "Unsupported layer_type: 'MonvConv'. " + "Supported types are: ARMAConv, GCSConv, GATConv, GCNConv, APPNPConv") # Test with invalud num_layers layer_type = "GATConv" num_layers = 0 @@ -195,7 +196,7 @@ def test_get_model(self): model = get_model(layer_type, num_layers, num_channels, activation, num_output) self.assertEqual(str( - error.exception), "num_layers must be at least 1.") + error.exception), "num_layers must be at least 1, got 0.") # Test with invalid num_output num_layers = 2 num_output = 0 @@ -203,7 +204,7 @@ def test_get_model(self): model = get_model(layer_type, num_layers, num_channels, activation, num_output) self.assertEqual(str( - error.exception), "num_output must be at least 1.") + error.exception), "num_output must be at least 1, got 0.") # Test with invalid num_channels num_output = 2 num_channels = 0 @@ -211,7 +212,7 @@ def test_get_model(self): model = get_model(layer_type, num_layers, num_channels, activation, num_output) self.assertEqual(str( - error.exception), "num_channels must be at least 1.") + error.exception), "num_channels must be at least 1, got 0.") # Test with invalid activation num_channels = 16 activation = 5 @@ -220,7 +221,7 @@ def test_get_model(self): num_channels, activation, num_output) self.assertEqual( str(error.exception), - "activation must be a string representing the activation function.") + "activation must be a string, got int.") def test_create_dataset(self): @@ -428,21 +429,17 @@ def test_perform_grid_search(self): num_layers = [1] num_channels = [8] activations = ["relu"] - model_parameters = [(layer, n_layer, channel, activation) - for layer in layers for n_layer in num_layers - for channel in num_channels - for activation in activations - for layer in layers for n_layer in num_layers] + parameter_grid = [(layer, n_layer, channel, activation) + for layer in layers for n_layer in num_layers + for channel in num_channels + for activation in activations] batch_size = 2 - loss_function = MeanAbsolutePercentageError() - optimizer = tf.keras.optimizers.Adam() es_patience = 5 max_epochs = 2 - training_parameters = [batch_size, loss_function, - optimizer, es_patience, max_epochs] # Perform grid search with explicit save_dir to avoid os.path.realpath issues - perform_grid_search(model_parameters, training_parameters, - dataset, save_dir=self.path) + perform_grid_search(dataset, parameter_grid, self.path, + batch_size=batch_size, max_epochs=max_epochs, + es_patience=es_patience) # Check if the results file is created results_file = os.path.join( @@ -451,8 +448,8 @@ def test_perform_grid_search(self): # Check if the results file has the expected number of rows df_results = pd.read_csv(results_file) - self.assertEqual(len(df_results), len(model_parameters)) - self.assertEqual(len(df_results.columns), 10) + self.assertEqual(len(df_results), len(parameter_grid)) + self.assertEqual(len(df_results.columns), 12) def test_scale_data_valid_data(self): """Test utils.scale_data with valid input and label data.""" From d3f6ca4bc1a2778a637ce79d2fe394010622513d Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 16:23:56 +0100 Subject: [PATCH 31/41] add rtd for gnn --- docs/source/python/m-surrogate.rst | 346 +++++++++++++++++- .../surrogatemodel/GNN/save_ground_truth.py | 54 --- 2 files changed, 343 insertions(+), 57 deletions(-) delete mode 100644 pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/save_ground_truth.py diff --git a/docs/source/python/m-surrogate.rst b/docs/source/python/m-surrogate.rst index 3ea9e9782d..9b80e790a0 100644 --- a/docs/source/python/m-surrogate.rst +++ b/docs/source/python/m-surrogate.rst @@ -163,7 +163,347 @@ The `grid_search.py` and `hyperparameter_tuning.py` modules provide tools for sy - Visualization of hyperparameter importance - Selection of optimal model configurations -SECIR Groups Model ------------------- -To be added... + +Graph Neural Network (GNN) Surrogate Models +-------------------------------------------- + +The Graph Neural Network (GNN) module provides advanced surrogate models that leverage spatial connectivity and age-stratified epidemiological dynamics. These models are designed for immediate and reliable pandemic response by combining mechanistic expert knowledge with machine learning efficiency. + +Overview and Scientific Foundation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The GNN surrogate models are based on the research presented in: + +|Graph_Neural_Network_Surrogates| + +The implementation leverages the mechanistic ODE-SECIR model (see :doc:`ODE-SECIR documentation <../models/ode_secir>`) as the underlying expert model, using Python bindings to the C++ backend for efficient simulation during data generation. + +Module Structure +~~~~~~~~~~~~~~~~ + +The GNN module is located in `pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN `_ and consists of: + +- **data_generation.py**: Generates training and evaluation data by simulating epidemiological scenarios with the mechanistic SECIR model +- **network_architectures.py**: Defines various GNN architectures (GCN, GAT, GIN) with configurable layers and preprocessing +- **evaluate_and_train.py**: Implements training and evaluation pipelines for GNN models +- **grid_search.py**: Provides hyperparameter optimization through systematic grid search +- **GNN_utils.py**: Contains utility functions for data preprocessing, graph construction, and population data handling + +Data Generation +~~~~~~~~~~~~~~~ + +The data generation process in ``data_generation.py`` creates graph-structured training data through mechanistic simulations: + +.. code-block:: python + + from memilio.surrogatemodel.GNN import data_generation + + # Generate training dataset + dataset = data_generation.generate_dataset( + num_runs=1000, # Number of simulation scenarios + num_days=30, # Simulation horizon + num_age_groups=6, # Age stratification + data_dir='path/to/contact_data', # Contact matrices location + mobility_dir='path/to/mobility', # Mobility data location + save_path='gnn_training_data.pickle' + ) + +**Data Generation Workflow:** + +1. **Parameter Sampling**: Randomly sample epidemiological parameters (transmission rates, incubation periods, recovery rates) from predefined distributions to create diverse scenarios. + +2. **Compartment Initialization**: Initialize epidemic compartments for each age group in each region based on realistic demographic data. Compartments are initialized using shared base factors. + +3. **Mobility Graph Construction**: Build a spatial graph where: + + - Nodes represent geographic regions (e.g., German counties) + - Edges represent mobility connections with weights from commuting data + - Node features include age-stratified population sizes + +4. **Contact Matrix Configuration**: Load and configure baseline contact matrices for different location types (home, school, work, other) stratified by age groups. + +5. **Damping Application**: Apply time-varying dampings to contact matrices to simulate NPIs: + + - Multiple damping periods with random start days + - Location-specific damping factors (e.g., stronger school closures, moderate workplace restrictions) + - Realistic parameter ranges based on observed intervention strengths + +6. **Simulation Execution**: Run the mechanistic ODE-SECIR model using MEmilio's C++ backend through Python bindings to generate the dataset. + +7. **Data Processing**: Transform simulation results into graph-structured format: + + - Extract compartment time series for each node (region) and age group + - Apply logarithmic transformation for numerical stability + - Store graph topology, node features, and temporal sequences + +Network Architectures +~~~~~~~~~~~~~~~~~~~~~ + +The ``network_architectures.py`` module provides flexible GNN model construction for different layer types. + +.. code-block:: python + + from memilio.surrogatemodel.GNN import network_architectures + + # Define GNN architecture + model_config = { + 'layer_type': 'GCN', # GNN layer type + 'num_layers': 3, # Network depth + 'hidden_dim': 64, # Hidden layer dimensions + 'activation': 'relu', # Activation function + 'dropout_rate': 0.2, # Dropout for regularization + 'use_batch_norm': True, # Batch normalization + 'aggregation': 'mean', # Neighborhood aggregation method + } + + # Build model + model = network_architectures.build_gnn_model( + config=model_config, + input_shape=(num_timesteps, num_features), + output_dim=num_compartments * num_age_groups + ) + + +Training and Evaluation +~~~~~~~~~~~~~~~~~~~~~~~ + +The ``evaluate_and_train.py`` module provides the training functionality: + +.. code-block:: python + + from memilio.surrogatemodel.GNN import evaluate_and_train + + # Load training data + with open('gnn_training_data.pickle', 'rb') as f: + dataset = pickle.load(f) + + # Define training configuration + training_config = { + 'epochs': 100, + 'batch_size': 32, + 'learning_rate': 0.001, + 'optimizer': 'adam', + 'loss_function': 'mse', + 'early_stopping_patience': 10, + 'validation_split': 0.2 + } + + # Train model + history = evaluate_and_train.train_gnn_model( + model=model, + dataset=dataset, + config=training_config, + save_weights='best_gnn_model.h5' + ) + + # Evaluate on test set + metrics = evaluate_and_train.evaluate_model( + model=model, + test_data=test_dataset, + metrics=['mae', 'mape', 'r2'] + ) + +**Training Features:** + +1. **Mini-batch Training**: Graph batching for efficient training on large datasets +2. **Custom Loss Functions**: MSE, MAE, MAPE, or custom compartment-weighted losses +3. **Early Stopping**: Monitors validation loss to prevent overfitting +4. **Learning Rate Scheduling**: Adaptive learning rate reduction on plateaus +5. **Save Best Weights**: Saves best model weights based on validation performance + +**Evaluation Metrics:** + +- **Mean Absolute Error (MAE)**: Average absolute prediction error per compartment +- **Mean Absolute Percentage Error (MAPE)**: Mean absolute error as percentage +- **R² Score**: Coefficient of determination for prediction quality + +**Data Splitting:** + +- **Training Set (70%)**: For model parameter optimization +- **Validation Set (15%)**: For hyperparameter tuning and early stopping +- **Test Set (15%)**: For final performance evaluation + +Hyperparameter Optimization +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``grid_search.py`` module enables systematic exploration of hyperparameter space: + +.. code-block:: python + + from memilio.surrogatemodel.GNN import grid_search + + # Define search space + param_grid = { + 'layer_type': ['GCN', 'GAT', 'GIN'], + 'num_layers': [2, 3, 4, 5], + 'hidden_dim': [32, 64, 128, 256], + 'learning_rate': [0.001, 0.0005, 0.0001], + 'dropout_rate': [0.0, 0.1, 0.2, 0.3], + 'batch_size': [16, 32, 64], + 'activation': ['relu', 'elu', 'tanh'] + } + + # Run grid search with cross-validation + results = grid_search.run_hyperparameter_search( + param_grid=param_grid, + data_path='gnn_training_data.pickle', + cv_folds=5, + metric='mae', + save_results='grid_search_results.csv' + ) + + # Analyze best configuration + best_config = grid_search.get_best_configuration(results) + print(f"Best configuration: {best_config}") + +Utility Functions +~~~~~~~~~~~~~~~~~ + +The ``GNN_utils.py`` module provides essential helper functions used throughout the GNN workflow: + +**Data Preprocessing:** + +.. code-block:: python + + from memilio.surrogatemodel.GNN import GNN_utils + + # Remove confirmed compartments (simplify model) + simplified_data = GNN_utils.remove_confirmed_compartments( + dataset_entries=dataset, + num_groups=6 + ) + + # Apply logarithmic scaling + scaled_data = GNN_utils.scale_data( + data=dataset, + method='log', + epsilon=1e-6 # Small constant to avoid log(0) + ) + + # Load population data + population = GNN_utils.load_population_data( + data_dir='path/to/demographics', + age_groups=[0, 5, 15, 35, 60, 80] + ) + +**Graph Construction:** + +.. code-block:: python + + # Create mobility graph from commuting data + graph = GNN_utils.create_mobility_graph( + mobility_dir='path/to/mobility', + num_regions=401, # German counties + county_ids=county_list, + models=models_per_region # SECIR models for each region + ) + + # Get baseline contact matrix + contact_matrix = GNN_utils.get_baseline_contact_matrix( + data_dir='path/to/contact_matrices' + ) + +Practical Usage Example +~~~~~~~~~~~~~~~~~~~~~~~ + +Here is a complete example workflow from data generation to model evaluation: + +.. code-block:: python + + import pickle + from pathlib import Path + from memilio.surrogatemodel.GNN import ( + data_generation, + network_architectures, + evaluate_and_train + ) + + # Step 1: Generate training data + print("Generating training data...") + dataset = data_generation.generate_dataset( + num_runs=5000, + num_days=30, + num_age_groups=6, + data_dir='/path/to/memilio/data/Germany', + mobility_dir='/path/to/mobility_data', + save_path='gnn_dataset_5000.pickle' + ) + + # Step 2: Define and build GNN model + print("Building GNN model...") + model_config = { + 'layer_type': 'GCN', + 'num_layers': 4, + 'hidden_dim': 128, + 'activation': 'relu', + 'dropout_rate': 0.2, + 'use_batch_norm': True + } + + model = network_architectures.build_gnn_model( + config=model_config, + input_shape=(1, 48), # 6 age groups × 8 compartments + output_dim=48 # Predict all compartments + ) + + # Step 3: Train the model + print("Training model...") + training_config = { + 'epochs': 200, + 'batch_size': 32, + 'learning_rate': 0.001, + 'optimizer': 'adam', + 'loss_function': 'mae', + 'early_stopping_patience': 20, + 'validation_split': 0.2 + } + + history = evaluate_and_train.train_gnn_model( + model=model, + dataset=dataset, + config=training_config, + save_weights='gnn_weights_best.h5' + ) + + # Step 4: Evaluate on test data + print("Evaluating model...") + test_metrics = evaluate_and_train.evaluate_model( + model=model, + test_data='gnn_test_data.pickle', + metrics=['mae', 'mape', 'r2'] + ) + + # Print results + print(f"Test MAE: {test_metrics['mae']:.4f}") + print(f"Test MAPE: {test_metrics['mape']:.2f}%") + print(f"Test R²: {test_metrics['r2']:.4f}") + + # Step 5: Make predictions on new scenarios + with open('new_scenario.pickle', 'rb') as f: + new_data = pickle.load(f) + + predictions = model.predict(new_data) + print(f"Predictions shape: {predictions.shape}") + +**GPU Acceleration:** + +- TensorFlow automatically uses GPU when available +- Spektral layers are optimized for GPU execution +- Training time can be heavily reduced with appropriate GPU hardware + +Additional Resources +~~~~~~~~~~~~~~~~~~~~ + +**Code and Examples:** + +- `GNN Module `_ +- `GNN README `_ +- `Test Scripts `_ + +**Related Documentation:** + +- :doc:`ODE-SECIR Model <../models/ode_secir>` +- :doc:`MEmilio Simulation Package ` +- :doc:`Python Bindings ` + diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/save_ground_truth.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/save_ground_truth.py deleted file mode 100644 index 3b8310bf8f..0000000000 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/save_ground_truth.py +++ /dev/null @@ -1,54 +0,0 @@ -import os -import numpy as np -import pandas as pd -import json -import matplotlib.pyplot as plt -import h5py - - -def read_case_data(): - df = pd.read_hdf( - "/localdata1/hege_mn/memilio/data/Germany/pydata/Results_rki.h5") - return df - - -def read_results_h5(path, group_key='Group1'): - with h5py.File(path, 'r') as f: - keys = list(f.keys()) - res = {} - for i, key in enumerate(keys): - group = f[key] - total = group[group_key][()] - # remove confirmed compartments - sum_inf_no_symp = np.sum(total[:, [2, 3]], axis=1) - sum_inf_symp = np.sum(total[:, [4, 5]], axis=1) - total[:, 2] = sum_inf_no_symp - total[:, 4] = sum_inf_symp - total = np.delete(total, [3, 5], axis=1) - res[key] = total - return res - - -def main(): - path_rki_h5 = "/localdata1/hege_mn/memilio/data/Germany/pydata/Results_rki.h5" - num_age_groups = 6 - all_age_data_list = [] - group_keys = ['Group1', 'Group2', 'Group3', 'Group4', 'Group5', 'Group6'] - for age in range(num_age_groups): - age_data_dict = read_results_h5(path_rki_h5, group_keys[age]) - age_data_np = np.array(list(age_data_dict.values())) - all_age_data_list.append(age_data_np) - # Combine age groups to get shape (num_nodes, timesteps, num_features=48) - all_age_data = np.stack(all_age_data_list, axis=-1) - print(f"Shape of all_age_data: {all_age_data.shape}") - num_nodes, timesteps, _, _ = all_age_data.shape - - # Save the ground truth data - output_path = "/localdata1/hege_mn/memilio/data/Germany/pydata/ground_truth_all_nodes.pickle" - with open(output_path, 'wb') as f: - pd.to_pickle(all_age_data, f) - print(f"Ground truth data saved to {output_path}") - - -if __name__ == "__main__": - main() From e706a614788b74ecc2ac1f46caef036594f002e9 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 5 Nov 2025 16:27:08 +0100 Subject: [PATCH 32/41] readme --- .../memilio/surrogatemodel/GNN/README.md | 35 ++++++++++++++++--- 1 file changed, 30 insertions(+), 5 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md index 6c1dd3e54d..007496b748 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md @@ -1,5 +1,30 @@ -GNN package -=============== -1) Proceed the real world data by running extrapolate_data_gnn.py. The data is saved as "ground_truth_all_nodes.pickle" in -the pydata folder -2) Generate data using data_generation.py. The model is initialized using the real world data + some noise. \ No newline at end of file +# Graph Neural Network (GNN) Surrogate Models + +This module implements Graph Neural Network-based surrogate models for epidemiological simulations, specifically designed to accelerate and enhance pandemic response modeling. + +## Overview + +The GNN surrogate models are based on the research presented in: + +**Schmidt A, Zunker H, Heinlein A, Kühn MJ. (2025).** *Graph Neural Network Surrogates to leverage Mechanistic Expert Knowledge towards Reliable and Immediate Pandemic Response*. Submitted for publication. +https://doi.org/10.48550/arXiv.2411.06500 + +This implementation leverages the underlying [ODE SECIR model](https://memilio.readthedocs.io/en/latest/cpp/models/osecir.html) and applies Graph Neural Networks to create fast, reliable surrogate models that can be used for immediate pandemic response scenarios. The models are stratified by age groups and incorporate spatial connectivity through mobility data. + +## Module Structure + +The GNN module consists of the following components: + +- **`data_generation.py`**: Generates training and evaluation data for GNN surrogate models by simulating epidemiological scenarios using the mechanistic SECIR model. Handles parameter sampling, compartment initialization, damping factors, and mobility connections between regions. + +- **`network_architectures.py`**: Defines various GNN architectures using different layer types (e.g., Graph Convolutional Networks, Graph Attention Networks). Provides functionality to configure network depth, width, activation functions, and preprocessing transformations. + +- **`evaluate_and_train.py`**: Implements training and evaluation pipelines for GNN surrogate models. Loads generated data, trains models with specified hyperparameters, evaluates performance metrics, and saves trained model weights. + +- **`grid_search.py`**: Provides hyperparameter optimization through systematic grid search over network architectures, training configurations, and model parameters to identify optimal GNN configurations for epidemiological forecasting. + +- **`GNN_utils.py`**: Contains utility functions for data preprocessing, mobility graph creation, population data loading, data scaling/transformation, and other helper functions used throughout the GNN workflow. + +## Documentation + +Comprehensive documentation for the GNN surrogate models, including tutorials and usage examples, is available in our [documentation](https://memilio.readthedocs.io/en/latest/python/m-surrogate.html). From d077cc13ef0fca14043b20da1f036277f1dbdd19 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:02:17 +0100 Subject: [PATCH 33/41] rm unused functions --- docs/source/python/m-surrogate.rst | 12 +--- .../memilio/surrogatemodel/GNN/GNN_utils.py | 66 ------------------- 2 files changed, 2 insertions(+), 76 deletions(-) diff --git a/docs/source/python/m-surrogate.rst b/docs/source/python/m-surrogate.rst index 9b80e790a0..c218de68e2 100644 --- a/docs/source/python/m-surrogate.rst +++ b/docs/source/python/m-surrogate.rst @@ -375,16 +375,9 @@ The ``GNN_utils.py`` module provides essential helper functions used throughout ) # Apply logarithmic scaling - scaled_data = GNN_utils.scale_data( + scaled_inputs, scaled_labels = GNN_utils.scale_data( data=dataset, - method='log', - epsilon=1e-6 # Small constant to avoid log(0) - ) - - # Load population data - population = GNN_utils.load_population_data( - data_dir='path/to/demographics', - age_groups=[0, 5, 15, 35, 60, 80] + transform=True ) **Graph Construction:** @@ -506,4 +499,3 @@ Additional Resources - :doc:`ODE-SECIR Model <../models/ode_secir>` - :doc:`MEmilio Simulation Package ` - :doc:`Python Bindings ` - diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py index e98d0ef99f..89d627cb16 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -26,16 +26,11 @@ """ import os -from typing import List, Tuple import numpy as np -import pandas as pd from sklearn.preprocessing import FunctionTransformer from memilio.simulation.osecir import ModelGraph, set_edges -from memilio.epidata import getDataIntoPandasDataFrame as gd -from memilio.epidata import transformMobilityData as tmd -from memilio.epidata import modifyDataframeSeries as mdfs # Default number of compartments in the ODE-SECIR model (without confirmed compartments) @@ -152,67 +147,6 @@ def create_mobility_graph(mobility_dir, num_regions, county_ids, models): return graph -def transform_mobility_data(data_dir): - """Updates mobility data to merge Eisenach and Wartburgkreis counties. - - :param data_dir: Root directory containing Germany mobility data. - :returns: Path to the updated mobility directory. - """ - # Get mobility data directory - mobility_dir = os.path.join(data_dir, 'Germany/mobility/') - - # Update mobility files by merging Eisenach and Wartburgkreis - tmd.updateMobility2022(mobility_dir, mobility_file='twitter_scaled_1252') - tmd.updateMobility2022( - mobility_dir, mobility_file='commuter_mobility_2022') - - return mobility_dir - - -def load_population_data(data_dir, age_groups=None): - """Loads population data for counties stratified by age groups. - - Reads county-level population data and aggregates it into specified age groups. - Default age groups follow the standard 6-group structure used in the ODE-SECIR model. - - :param data_dir: Root directory containing population data (should contain Germany/pydata/ subdirectory). - :param age_groups: List of age group labels (default: ['0-4', '5-14', '15-34', '35-59', '60-79', '80-130']). - :returns: List of population data entries, where each entry is [county_id, pop_group1, pop_group2, ...]. - :raises FileNotFoundError: If population data file is not found. - """ - if age_groups is None: - age_groups = ['0-4', '5-14', '15-34', '35-59', '60-79', '80-130'] - - # Load population data - population_file = os.path.join( - data_dir, 'Germany', 'pydata', 'county_population.json') - - if not os.path.exists(population_file): - raise FileNotFoundError( - f"Population data file not found: {population_file}") - - df_population = pd.read_json(population_file) - - # Initialize DataFrame for age-grouped population - df_population_agegroups = pd.DataFrame( - columns=[df_population.columns[0]] + age_groups - ) - - # Process each region - for region_id in df_population.iloc[:, 0]: - region_data = df_population[df_population.iloc[:, 0] - == int(region_id)] - age_grouped_pop = list( - mdfs.fit_age_group_intervals( - region_data.iloc[:, 2:], age_groups) - ) - - df_population_agegroups.loc[len(df_population_agegroups.index)] = [ - int(region_id)] + age_grouped_pop - - return df_population_agegroups.values.tolist() - - def scale_data( data, transform=True, num_compartments=DEFAULT_NUM_COMPARTMENTS): """Applies logarithmic transformation to simulation data for training. From a6371d0ac110adf480f5979517b74f7ad81fbf99 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:08:00 +0100 Subject: [PATCH 34/41] . --- .../memilio/surrogatemodel/GNN/GNN_utils.py | 3 ++- .../memilio/surrogatemodel/GNN/README.md | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py index 89d627cb16..2979ff261e 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -57,7 +57,8 @@ def remove_confirmed_compartments(dataset_entries, num_groups): for timestep_data in dataset_entries: # Reshape to separate age groups and compartments data_reshaped = timestep_data.reshape( - [num_groups, int(np.asarray(dataset_entries).shape[1] / num_groups)] + [num_groups, int(np.asarray( + dataset_entries).shape[1] / num_groups)] ) # Merge InfectedNoSymptoms (index 2) with InfectedNoSymptomsConfirmed (index 3) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md index 007496b748..8ce7e823ce 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/README.md @@ -23,7 +23,7 @@ The GNN module consists of the following components: - **`grid_search.py`**: Provides hyperparameter optimization through systematic grid search over network architectures, training configurations, and model parameters to identify optimal GNN configurations for epidemiological forecasting. -- **`GNN_utils.py`**: Contains utility functions for data preprocessing, mobility graph creation, population data loading, data scaling/transformation, and other helper functions used throughout the GNN workflow. +- **`GNN_utils.py`**: Contains utility functions for data preprocessing (e.g., removing confirmed compartments, scaling data), and building graphs or baseline contact matrices used throughout the GNN workflow. ## Documentation From ec9c932f99f4f0fe12c472bfbfa5f4bec4d9dd66 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 26 Nov 2025 15:08:30 +0100 Subject: [PATCH 35/41] [ci skip] fix data_generation --- .../memilio/surrogatemodel/GNN/GNN_utils.py | 2 +- .../memilio/surrogatemodel/GNN/data_generation.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py index 2979ff261e..cea67fc304 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -176,7 +176,7 @@ def scale_data( raise ValueError("Label data must be numeric.") # Calculate number of age groups from data shape - num_groups = int(inputs_array.shape[2] / num_compartments) + num_groups = int(inputs_array.shape[-1] / num_compartments) # Initialize transformer (log1p for numerical stability) transformer = FunctionTransformer(np.log1p, validate=True) diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index eb3e44e271..3f7c84f218 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -40,7 +40,7 @@ import memilio.simulation.osecir as osecir from memilio.simulation import AgeGroup, set_log_level, Damping from memilio.simulation.osecir import ( - Index_InfectionState, InfectionState, ParameterStudy, + Index_InfectionState, InfectionState, GraphParameterStudy, interpolate_simulation_result) from memilio.surrogatemodel.GNN.GNN_utils import ( @@ -394,7 +394,7 @@ def run_secir_groups_simulation( graph.get_node(node_idx).property.populations = model.populations # Run simulation and measure runtime - study = ParameterStudy(graph, t0=0, tmax=days, dt=0.5, num_runs=1) + study = GraphParameterStudy(graph, 0, days, 0.5, 1) start_time = time.perf_counter() study_results = study.run() From 1b90974b085446182d30b6e50eb9f0e2b79ca227 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Wed, 26 Nov 2025 15:23:07 +0100 Subject: [PATCH 36/41] rework some files --- docs/source/python/m-surrogate.rst | 264 ++++++++---------- .../memilio/surrogatemodel/GNN/GNN_utils.py | 56 +--- .../surrogatemodel/GNN/data_generation.py | 27 +- .../surrogatemodel/GNN/evaluate_and_train.py | 32 ++- 4 files changed, 169 insertions(+), 210 deletions(-) diff --git a/docs/source/python/m-surrogate.rst b/docs/source/python/m-surrogate.rst index c218de68e2..6edc20aeb8 100644 --- a/docs/source/python/m-surrogate.rst +++ b/docs/source/python/m-surrogate.rst @@ -185,7 +185,7 @@ Module Structure The GNN module is located in `pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN `_ and consists of: - **data_generation.py**: Generates training and evaluation data by simulating epidemiological scenarios with the mechanistic SECIR model -- **network_architectures.py**: Defines various GNN architectures (GCN, GAT, GIN) with configurable layers and preprocessing +- **network_architectures.py**: Defines various GNN architectures (ARMAConv, GCSConv, GATConv, GCNConv, APPNPConv) with configurable depth and channels - **evaluate_and_train.py**: Implements training and evaluation pipelines for GNN models - **grid_search.py**: Provides hyperparameter optimization through systematic grid search - **GNN_utils.py**: Contains utility functions for data preprocessing, graph construction, and population data handling @@ -193,20 +193,24 @@ The GNN module is located in `pycode/memilio-surrogatemodel/memilio/surrogatemod Data Generation ~~~~~~~~~~~~~~~ -The data generation process in ``data_generation.py`` creates graph-structured training data through mechanistic simulations: +The data generation process in ``data_generation.py`` creates graph-structured training data through mechanistic simulations. Use ``generate_data`` to run multiple simulations and persist a pickle with inputs, labels, damping info, and contact matrices: .. code-block:: python from memilio.surrogatemodel.GNN import data_generation - - # Generate training dataset - dataset = data_generation.generate_dataset( - num_runs=1000, # Number of simulation scenarios - num_days=30, # Simulation horizon - num_age_groups=6, # Age stratification - data_dir='path/to/contact_data', # Contact matrices location - mobility_dir='path/to/mobility', # Mobility data location - save_path='gnn_training_data.pickle' + import memilio.simulation as mio + + data = data_generation.generate_data( + num_runs=5, + data_dir="/path/to/memilio/data", + output_path="/tmp/generated_datasets", + input_width=5, + label_width=30, + start_date=mio.Date(2020, 10, 1), + end_date=mio.Date(2021, 10, 31), + mobility_file="commuter_mobility.txt", # or commuter_mobility_2022.txt + transform=True, + save_data=True ) **Data Generation Workflow:** @@ -240,28 +244,18 @@ The data generation process in ``data_generation.py`` creates graph-structured t Network Architectures ~~~~~~~~~~~~~~~~~~~~~ -The ``network_architectures.py`` module provides flexible GNN model construction for different layer types. +The ``network_architectures.py`` module provides flexible GNN model construction for supported layer types (ARMAConv, GCSConv, GATConv, GCNConv, APPNPConv). .. code-block:: python from memilio.surrogatemodel.GNN import network_architectures - - # Define GNN architecture - model_config = { - 'layer_type': 'GCN', # GNN layer type - 'num_layers': 3, # Network depth - 'hidden_dim': 64, # Hidden layer dimensions - 'activation': 'relu', # Activation function - 'dropout_rate': 0.2, # Dropout for regularization - 'use_batch_norm': True, # Batch normalization - 'aggregation': 'mean', # Neighborhood aggregation method - } - - # Build model - model = network_architectures.build_gnn_model( - config=model_config, - input_shape=(num_timesteps, num_features), - output_dim=num_compartments * num_age_groups + + model = network_architectures.get_model( + layer_type="GCNConv", + num_layers=3, + num_channels=64, + activation="relu", + num_output=48 # outputs per node ) @@ -272,36 +266,34 @@ The ``evaluate_and_train.py`` module provides the training functionality: .. code-block:: python - from memilio.surrogatemodel.GNN import evaluate_and_train - - # Load training data - with open('gnn_training_data.pickle', 'rb') as f: - dataset = pickle.load(f) - - # Define training configuration - training_config = { - 'epochs': 100, - 'batch_size': 32, - 'learning_rate': 0.001, - 'optimizer': 'adam', - 'loss_function': 'mse', - 'early_stopping_patience': 10, - 'validation_split': 0.2 - } - - # Train model - history = evaluate_and_train.train_gnn_model( - model=model, - dataset=dataset, - config=training_config, - save_weights='best_gnn_model.h5' + from tensorflow.keras.losses import MeanAbsolutePercentageError + from tensorflow.keras.optimizers import Adam + from memilio.surrogatemodel.GNN import evaluate_and_train, network_architectures + + dataset = evaluate_and_train.load_gnn_dataset( + "/tmp/generated_datasets/GNN_data_30days_3dampings_classic5.pickle", + "/path/to/memilio/data/Germany/mobility", + number_of_nodes=400 ) - - # Evaluate on test set - metrics = evaluate_and_train.evaluate_model( + + model = network_architectures.get_model( + layer_type="GCNConv", + num_layers=3, + num_channels=32, + activation="relu", + num_output=48 + ) + + results = evaluate_and_train.train_and_evaluate( + data=dataset, + batch_size=32, + epochs=50, model=model, - test_data=test_dataset, - metrics=['mae', 'mape', 'r2'] + loss_fn=MeanAbsolutePercentageError(), + optimizer=Adam(learning_rate=0.001), + es_patience=10, + save_dir="/tmp/model_results", + save_name="gnn_model" ) **Training Features:** @@ -309,8 +301,7 @@ The ``evaluate_and_train.py`` module provides the training functionality: 1. **Mini-batch Training**: Graph batching for efficient training on large datasets 2. **Custom Loss Functions**: MSE, MAE, MAPE, or custom compartment-weighted losses 3. **Early Stopping**: Monitors validation loss to prevent overfitting -4. **Learning Rate Scheduling**: Adaptive learning rate reduction on plateaus -5. **Save Best Weights**: Saves best model weights based on validation performance +4. **Save Best Weights**: Saves best model weights based on validation performance **Evaluation Metrics:** @@ -331,31 +322,31 @@ The ``grid_search.py`` module enables systematic exploration of hyperparameter s .. code-block:: python - from memilio.surrogatemodel.GNN import grid_search - - # Define search space - param_grid = { - 'layer_type': ['GCN', 'GAT', 'GIN'], - 'num_layers': [2, 3, 4, 5], - 'hidden_dim': [32, 64, 128, 256], - 'learning_rate': [0.001, 0.0005, 0.0001], - 'dropout_rate': [0.0, 0.1, 0.2, 0.3], - 'batch_size': [16, 32, 64], - 'activation': ['relu', 'elu', 'tanh'] - } - - # Run grid search with cross-validation - results = grid_search.run_hyperparameter_search( - param_grid=param_grid, - data_path='gnn_training_data.pickle', - cv_folds=5, - metric='mae', - save_results='grid_search_results.csv' + from pathlib import Path + from memilio.surrogatemodel.GNN import grid_search, evaluate_and_train + + data = evaluate_and_train.create_dataset( + "/tmp/generated_datasets/GNN_data_30days_3dampings_classic5.pickle", + "/path/to/memilio/data/Germany/mobility", + number_of_nodes=400 + ) + + parameter_grid = grid_search.generate_parameter_grid( + layer_types=["GCNConv", "GATConv"], + num_layers_options=[2, 3], + num_channels_options=[16, 32], + activation_functions=["relu", "elu"] + ) + + grid_search.perform_grid_search( + data=data, + parameter_grid=parameter_grid, + save_dir=str(Path("/tmp/grid_results")), + batch_size=32, + max_epochs=50, + es_patience=10, + learning_rate=0.001 ) - - # Analyze best configuration - best_config = grid_search.get_best_configuration(results) - print(f"Best configuration: {best_config}") Utility Functions ~~~~~~~~~~~~~~~~~ @@ -404,80 +395,55 @@ Here is a complete example workflow from data generation to model evaluation: .. code-block:: python - import pickle - from pathlib import Path + import memilio.simulation as mio + from tensorflow.keras.losses import MeanAbsolutePercentageError + from tensorflow.keras.optimizers import Adam from memilio.surrogatemodel.GNN import ( - data_generation, - network_architectures, + data_generation, + network_architectures, evaluate_and_train ) - - # Step 1: Generate training data - print("Generating training data...") - dataset = data_generation.generate_dataset( - num_runs=5000, - num_days=30, - num_age_groups=6, - data_dir='/path/to/memilio/data/Germany', - mobility_dir='/path/to/mobility_data', - save_path='gnn_dataset_5000.pickle' + + # Step 1: Generate and save training data + data_generation.generate_data( + num_runs=100, + data_dir="/path/to/memilio/data", + output_path="/tmp/generated_datasets", + input_width=5, + label_width=30, + start_date=mio.Date(2020, 10, 1), + end_date=mio.Date(2021, 10, 31), + save_data=True, + mobility_file="commuter_mobility.txt" ) - - # Step 2: Define and build GNN model - print("Building GNN model...") - model_config = { - 'layer_type': 'GCN', - 'num_layers': 4, - 'hidden_dim': 128, - 'activation': 'relu', - 'dropout_rate': 0.2, - 'use_batch_norm': True - } - - model = network_architectures.build_gnn_model( - config=model_config, - input_shape=(1, 48), # 6 age groups × 8 compartments - output_dim=48 # Predict all compartments + + # Step 2: Load dataset and build model + dataset = evaluate_and_train.load_gnn_dataset( + "/tmp/generated_datasets/GNN_data_30days_3dampings_classic100.pickle", + "/path/to/memilio/data/Germany/mobility", + number_of_nodes=400 ) - - # Step 3: Train the model - print("Training model...") - training_config = { - 'epochs': 200, - 'batch_size': 32, - 'learning_rate': 0.001, - 'optimizer': 'adam', - 'loss_function': 'mae', - 'early_stopping_patience': 20, - 'validation_split': 0.2 - } - - history = evaluate_and_train.train_gnn_model( - model=model, - dataset=dataset, - config=training_config, - save_weights='gnn_weights_best.h5' + + model = network_architectures.get_model( + layer_type="GCNConv", + num_layers=4, + num_channels=128, + activation="relu", + num_output=48 ) - - # Step 4: Evaluate on test data - print("Evaluating model...") - test_metrics = evaluate_and_train.evaluate_model( + + # Step 3: Train and evaluate + results = evaluate_and_train.train_and_evaluate( + data=dataset, + batch_size=32, + epochs=100, model=model, - test_data='gnn_test_data.pickle', - metrics=['mae', 'mape', 'r2'] + loss_fn=MeanAbsolutePercentageError(), + optimizer=Adam(learning_rate=0.001), + es_patience=20, + save_dir="/tmp/model_results", + save_name="gnn_weights_best" ) - - # Print results - print(f"Test MAE: {test_metrics['mae']:.4f}") - print(f"Test MAPE: {test_metrics['mape']:.2f}%") - print(f"Test R²: {test_metrics['r2']:.4f}") - - # Step 5: Make predictions on new scenarios - with open('new_scenario.pickle', 'rb') as f: - new_data = pickle.load(f) - - predictions = model.predict(new_data) - print(f"Predictions shape: {predictions.shape}") **GPU Acceleration:** diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py index cea67fc304..4d4e074855 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/GNN_utils.py @@ -175,56 +175,22 @@ def scale_data( if not np.issubdtype(labels_array.dtype, np.number): raise ValueError("Label data must be numeric.") - # Calculate number of age groups from data shape - num_groups = int(inputs_array.shape[-1] / num_compartments) - # Initialize transformer (log1p for numerical stability) transformer = FunctionTransformer(np.log1p, validate=True) - # Process inputs - # Reshape: [samples, timesteps, nodes, features] -> [nodes, samples, timesteps, features] - # -> [nodes * compartments, samples * timesteps * age_groups] - inputs_reshaped = inputs_array.transpose(2, 0, 1, 3).reshape( - num_groups * num_compartments, -1 - ) - if transform: - inputs_transformed = transformer.transform(inputs_reshaped) + inputs_scaled = transformer.transform( + inputs_array.reshape(-1, inputs_array.shape[-1]) + ).reshape(inputs_array.shape) + labels_scaled = transformer.transform( + labels_array.reshape(-1, labels_array.shape[-1]) + ).reshape(labels_array.shape) else: - inputs_transformed = inputs_reshaped - - original_shape_input = inputs_array.shape - - # Reverse reshape to separate dimensions - inputs_back = inputs_transformed.reshape( - original_shape_input[2], - original_shape_input[0], - original_shape_input[1], - original_shape_input[3] - ) - - # Reverse transpose and reorder to [samples, features, timesteps, nodes] - scaled_inputs = inputs_back.transpose(1, 2, 0, 3).transpose(0, 3, 1, 2) - - # Process labels with same procedure - labels_reshaped = labels_array.transpose(2, 0, 1, 3).reshape( - num_groups * num_compartments, -1 - ) - - if transform: - labels_transformed = transformer.transform(labels_reshaped) - else: - labels_transformed = labels_reshaped - - original_shape_labels = labels_array.shape - - labels_back = labels_transformed.reshape( - original_shape_labels[2], - original_shape_labels[0], - original_shape_labels[1], - original_shape_labels[3] - ) + inputs_scaled = inputs_array + labels_scaled = labels_array - scaled_labels = labels_back.transpose(1, 2, 0, 3).transpose(0, 3, 1, 2) + # Reorder to [samples, features, timesteps, nodes] + scaled_inputs = inputs_scaled.transpose(0, 3, 1, 2) + scaled_labels = labels_scaled.transpose(0, 3, 1, 2) return scaled_inputs, scaled_labels diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py index 3f7c84f218..51873845bc 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/data_generation.py @@ -469,9 +469,9 @@ def generate_data(num_runs, data_dir, output_path, input_width, label_width, if not os.path.exists(mobility_path): raise FileNotFoundError(f"Mobility file not found: {mobility_path}") - # Create graph (reused for all runs with different initial conditions) - graph = get_graph(num_groups, data_dir, - mobility_path, start_date, end_date) + # Create base graph. A new copy will be used per run to avoid cumulative dampings. + base_graph = get_graph(num_groups, data_dir, + mobility_path, start_date, end_date) print(f"\nGenerating {num_runs} simulation runs...") bar = Bar( @@ -495,11 +495,13 @@ def generate_data(num_runs, data_dir, output_path, input_width, label_width, damping_factors = [] # Run simulation - simulation_result, damped_mats, damping_coeffs, runtime = \ - run_secir_groups_simulation( - total_days, damping_days, damping_factors, graph, - within_group_variation, num_groups - ) + # Use a fresh graph copy per run to prevent damping accumulation across runs. + graph = copy.deepcopy(base_graph) + + simulation_result, damped_mats, damping_coeffs, runtime = run_secir_groups_simulation( + total_days, damping_days, damping_factors, graph, + within_group_variation, num_groups + ) runtimes.append(runtime) @@ -529,12 +531,15 @@ def generate_data(num_runs, data_dir, output_path, input_width, label_width, # Save dataset if requested if save_data: - # Apply scaling transformation + # Apply scaling transformation and store with shape + # [samples, timesteps, nodes, features] expected by loaders. inputs_scaled, labels_scaled = scale_data(data, transform) + inputs_for_save = inputs_scaled.transpose(0, 2, 3, 1) + labels_for_save = labels_scaled.transpose(0, 2, 3, 1) all_data = { - "inputs": inputs_scaled, - "labels": labels_scaled, + "inputs": inputs_for_save, + "labels": labels_for_save, "damping_day": data["damping_days"], "contact_matrix": data["contact_matrix"], "damping_coeff": data["damping_factors"] diff --git a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py index 75e8fce937..13ac944d6c 100644 --- a/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py +++ b/pycode/memilio-surrogatemodel/memilio/surrogatemodel/GNN/evaluate_and_train.py @@ -119,9 +119,26 @@ def load_gnn_dataset( if not dataset_path.exists(): raise FileNotFoundError(f"Dataset not found: {dataset_path}") - mobility_path = mobility_dir / mobility_filename - if not mobility_path.exists(): - raise FileNotFoundError(f"Mobility file not found: {mobility_path}") + # Accept multiple common commuter mobility filenames for compatibility + candidate_files = [] + if mobility_filename: + candidate_files.append(mobility_filename) + candidate_files.extend([ + "commuter_mobility.txt", + "commuter_mobility_2022.txt" + ]) + mobility_path = None + for fname in candidate_files: + path = mobility_dir / fname + if path.exists(): + mobility_path = path + break + + if mobility_path is None: + raise FileNotFoundError( + f"Mobility file not found in {mobility_dir}. " + f"Tried: {', '.join(candidate_files)}" + ) with dataset_path.open("rb") as fp: data = pickle.load(fp) @@ -137,11 +154,16 @@ def load_gnn_dataset( raise ValueError( "Loaded dataset is empty; expected at least one sample.") + if inputs.ndim != 4 or labels.ndim != 4: + raise ValueError( + f"Expected 4D arrays for inputs/labels, got {inputs.ndim} and {labels.ndim}." + ) + # Flatten temporal dimensions into feature vectors per node. num_samples, input_width, num_nodes, num_features = inputs.shape - _, label_width, _, label_features = labels.shape + _, label_width, label_nodes, label_features = labels.shape - if num_nodes != number_of_nodes: + if num_nodes != number_of_nodes or label_nodes != number_of_nodes: raise ValueError( f"Number of nodes in dataset ({num_nodes}) does not match expected " f"value ({number_of_nodes}).") From b24e203dc4044c2ba9906c311d0ee30c48bde775 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Thu, 27 Nov 2025 08:32:12 +0100 Subject: [PATCH 37/41] fix in doc --- docs/source/python/m-surrogate.rst | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/python/m-surrogate.rst b/docs/source/python/m-surrogate.rst index 6edc20aeb8..6ef53f177f 100644 --- a/docs/source/python/m-surrogate.rst +++ b/docs/source/python/m-surrogate.rst @@ -1,3 +1,5 @@ +.. include:: ../literature.rst + MEmilio Surrogate Model ======================== @@ -5,7 +7,7 @@ MEmilio Surrogate Model contains machine learning based surrogate models that ma Currently there are only surrogate models for ODE-type models. The simulations of these models are used for data generation. The goal is to create a powerful tool that predicts the infection dynamics faster than a simulation of an expert model, e.g., a metapopulation or agent-based model while still having acceptable errors with respect to the original simulations. - + The package can be found in `pycode/memilio-surrogatemodel `_. For more details, we refer to: Schmidt A, Zunker H, Heinlein A, Kühn MJ. (2025). *Graph Neural Network Surrogates to leverage Mechanistic Expert Knowledge towards Reliable and Immediate Pandemic Response*. Submitted for publication. `arXiv:2411.06500 `_ @@ -177,7 +179,7 @@ The GNN surrogate models are based on the research presented in: |Graph_Neural_Network_Surrogates| -The implementation leverages the mechanistic ODE-SECIR model (see :doc:`ODE-SECIR documentation <../models/ode_secir>`) as the underlying expert model, using Python bindings to the C++ backend for efficient simulation during data generation. +The implementation leverages the mechanistic ODE-SECIR model (see :doc:`ODE-SECIR documentation `) as the underlying expert model, using Python bindings to the C++ backend for efficient simulation during data generation. Module Structure ~~~~~~~~~~~~~~~~ @@ -458,10 +460,8 @@ Additional Resources - `GNN Module `_ - `GNN README `_ -- `Test Scripts `_ **Related Documentation:** -- :doc:`ODE-SECIR Model <../models/ode_secir>` +- :doc:`ODE-SECIR Model ` - :doc:`MEmilio Simulation Package ` -- :doc:`Python Bindings ` From f09f38a4fa516ed3a286d1a97161edbe94b92503 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Thu, 27 Nov 2025 08:33:14 +0100 Subject: [PATCH 38/41] more --- docs/source/python/m-surrogate.rst | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/python/m-surrogate.rst b/docs/source/python/m-surrogate.rst index 6ef53f177f..65386f592d 100644 --- a/docs/source/python/m-surrogate.rst +++ b/docs/source/python/m-surrogate.rst @@ -10,7 +10,9 @@ e.g., a metapopulation or agent-based model while still having acceptable errors The package can be found in `pycode/memilio-surrogatemodel `_. -For more details, we refer to: Schmidt A, Zunker H, Heinlein A, Kühn MJ. (2025). *Graph Neural Network Surrogates to leverage Mechanistic Expert Knowledge towards Reliable and Immediate Pandemic Response*. Submitted for publication. `arXiv:2411.06500 `_ +For more details, we refer to: + +|Graph_Neural_Network_Surrogates| Dependencies ------------ From 7d3ff083ece58b0c4d33ca975efb91d3430b17d2 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Thu, 27 Nov 2025 09:42:13 +0100 Subject: [PATCH 39/41] formatting doc --- docs/source/python/m-surrogate.rst | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/docs/source/python/m-surrogate.rst b/docs/source/python/m-surrogate.rst index 65386f592d..3104a4af57 100644 --- a/docs/source/python/m-surrogate.rst +++ b/docs/source/python/m-surrogate.rst @@ -34,17 +34,19 @@ Usage The package currently provides the following modules: - `models`: models for different tasks - Currently we have the following models: - - `ode_secir_simple`: A simple model allowing for asymptomatic as well as symptomatic infection states not stratified by age groups. - - `ode_secir_groups`: A model allowing for asymptomatic as well as symptomatic infection states stratified by age groups and including one damping. - Each model folder contains the following files: - - `data_generation`: Data generation from expert model simulation outputs. - - `model`: Training and evaluation of the model. - - `network_architectures`: Contains multiple network architectures. - - `grid_search`: Utilities for hyperparameter optimization. - - `hyperparameter_tuning`: Scripts for tuning model hyperparameters. + Currently we have the following models: + - `ode_secir_simple`: A simple model allowing for asymptomatic as well as symptomatic infection states not stratified by age groups. + - `ode_secir_groups`: A model allowing for asymptomatic as well as symptomatic infection states stratified by age groups and including one damping. + + Each model folder contains the following files: + + - `data_generation`: Data generation from expert model simulation outputs. + - `model`: Training and evaluation of the model. + - `network_architectures`: Contains multiple network architectures. + - `grid_search`: Utilities for hyperparameter optimization. + - `hyperparameter_tuning`: Scripts for tuning model hyperparameters. - `tests`: This file contains all tests. @@ -465,5 +467,4 @@ Additional Resources **Related Documentation:** -- :doc:`ODE-SECIR Model ` - :doc:`MEmilio Simulation Package ` From 141a2eeb7610bec07a74fe09a59ab2017f302614 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Thu, 4 Dec 2025 13:36:24 +0100 Subject: [PATCH 40/41] add deps to surrogate toml --- pycode/memilio-surrogatemodel/pyproject.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pycode/memilio-surrogatemodel/pyproject.toml b/pycode/memilio-surrogatemodel/pyproject.toml index 2660f0139a..d4451ae9d4 100644 --- a/pycode/memilio-surrogatemodel/pyproject.toml +++ b/pycode/memilio-surrogatemodel/pyproject.toml @@ -21,7 +21,9 @@ dependencies = [ "matplotlib", "scikit-learn", # first support of python 3.11 - "pyfakefs>=4.6" + "pyfakefs>=4.6", + "spektral>=1.2", + "tensorflow>=2.12.0" ] [project.optional-dependencies] From b5b360521019cccb9a661e7528bca366dc2ca988 Mon Sep 17 00:00:00 2001 From: HenrZu <69154294+HenrZu@users.noreply.github.com> Date: Fri, 5 Dec 2025 08:39:22 +0100 Subject: [PATCH 41/41] test file was lost in merge --- .../tests/test_surrogatemodel_GNN.py | 488 ++++++++++++++++++ 1 file changed, 488 insertions(+) create mode 100644 pycode/memilio-surrogatemodel/tests/test_surrogatemodel_GNN.py diff --git a/pycode/memilio-surrogatemodel/tests/test_surrogatemodel_GNN.py b/pycode/memilio-surrogatemodel/tests/test_surrogatemodel_GNN.py new file mode 100644 index 0000000000..58419294fb --- /dev/null +++ b/pycode/memilio-surrogatemodel/tests/test_surrogatemodel_GNN.py @@ -0,0 +1,488 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Manuel Heger, Henrik Zunker +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + +from pyfakefs import fake_filesystem_unittest +from unittest.mock import patch +import os +import unittest +import pickle +import numpy as np +import pandas as pd +import tensorflow as tf +import logging +import spektral + +import memilio.surrogatemodel.GNN.network_architectures as gnn_arch +import memilio.surrogatemodel.GNN.GNN_utils as utils + +from memilio.surrogatemodel.GNN.evaluate_and_train import ( + create_dataset, train_and_evaluate, evaluate, train_step, MixedLoader) +from memilio.surrogatemodel.GNN.grid_search import perform_grid_search +from memilio.surrogatemodel.GNN.network_architectures import generate_model_class, get_model +from tensorflow.keras.losses import MeanAbsolutePercentageError +# suppress all autograph warnings from tensorflow + +logging.getLogger("tensorflow").setLevel(logging.ERROR) + + +class TestSurrogatemodelGNN(fake_filesystem_unittest.TestCase): + path = "/home/" + + def create_dummy_data( + self, num_samples, num_nodes, num_node_features, output_dim): + """ + Create dummy data for testing. + + :param num_samples: Number of samples in the dataset. + :param num_nodes: Number of nodes in each graph. + :param num_node_features: Number of features per node. + :param output_dim: Number of output dimensions per node. + :return: A dictionary containing inputs, adjacency matrix, and labels. + """ + # Shape should be (num_samples, input_width, num_nodes, num_features) + X = np.random.rand(num_samples, 1, num_nodes, + num_node_features).astype(np.float32) + A = np.random.randint(0, 2, (num_nodes, num_nodes)).astype(np.float32) + # Shape should be (num_samples, label_width, num_nodes, label_features) + y = np.random.rand(num_samples, 1, num_nodes, + output_dim).astype(np.float32) + return {"inputs": X, "adjacency": A, "labels": y} + + def setup_fake_filesystem(self, fs, path, data): + """ + Save dummy data to the fake file system. + + :param fs: The fake file system object. + :param path: The base path in the fake file system. + :param data: The dummy data dictionary containing inputs, adjacency, and labels. + :return: Paths to the cases and mobility files. + """ + path_cases_dir = os.path.join(path, "cases") + fs.create_dir(path_cases_dir) + path_cases = os.path.join(path_cases_dir, "cases.pickle") + with open(path_cases, 'wb') as f: + pickle.dump({"inputs": data["inputs"], + "labels": data["labels"]}, f) + + path_mobility = os.path.join(path, "mobility") + mobility_file = os.path.join( + path_mobility, "commuter_mobility_2022.txt") + fs.create_dir(path_mobility) + fs.create_file(mobility_file) + with open(mobility_file, 'w') as f: + np.savetxt(f, data["adjacency"], delimiter=" ") + + return path_cases, path_mobility + + def setUp(self): + self.setUpPyfakefs() + + def test_generate_model_class(self): + + # Test parameters + layer_types = [ + lambda: tf.keras.layers.Dense(10, activation="relu"), + ] + num_layers = [3] + num_output = 2 + + # Generate the model class + ModelClass = generate_model_class( + "TestModel", layer_types, num_layers, num_output) + # Check if the generated class is a subclass of tf.keras.Model + self.assertTrue(issubclass(ModelClass, tf.keras.Model)) + + # Instantiate the model + model = ModelClass() + + # Check if the model has the expected number of layers + expected_num_layers = num_layers[0] + 1 # +1 for the output layer + self.assertEqual(len(model.layers), expected_num_layers) + self.assertIsInstance(model.layers[-1], tf.keras.layers.Dense) + self.assertEqual(model.layers[-1].units, num_output) + self.assertEqual( + model.layers[-1].activation.__name__, "relu") + + # Test with invalid parameters + layer_types = [ + lambda: tf.keras.layers.Dense(10, activation="relu"), + ] + num_layers = [0] + num_output = 2 + with self.assertRaises(ValueError) as error: + ModelClass = generate_model_class( + "TestModel", layer_types, num_layers, num_output) + model = ModelClass() + self.assertEqual(str( + error.exception), "All values in num_repeat must be at least 1.") + + num_layers = [3, 2] + with self.assertRaises(ValueError) as error: + ModelClass = generate_model_class( + "TestModel", layer_types, num_layers, num_output) + model = ModelClass() + self.assertEqual( + str(error.exception), + "layer_types and num_repeat must have the same length. " + "Got 1 and 2.") + + # Test with multiple layer types + layer_types = [ + lambda: tf.keras.layers.Dense(10, activation="relu"), + lambda: tf.keras.layers.Dense(20, activation="relu"), + ] + num_repeat = [2, 3] + num_output = 4 + + ModelClass = generate_model_class( + "TestModel", layer_types, num_repeat, num_output) + model = ModelClass() + + # Check the number of layers + self.assertEqual(len(model.layer_seq), sum(num_repeat)) + self.assertEqual(model.output_layer.units, num_output) + + def test_get_model(self): + + # Test parameters + layer_type = "GCNConv" + num_layers = 2 + num_channels = 16 + activation = "relu" + num_output = 3 + + # Generate the model + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + + # Check if the model is an instance of tf.keras.Model + self.assertIsInstance(model, tf.keras.Model) + + # Check if the model has the expected number of layers + expected_num_layers = num_layers + 1 # +1 for the output layer + self.assertEqual(len(model.layers), expected_num_layers) + + # Check handling of invalid parameters + # Test with invalid layer type + layer_type = "MonvConv" + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual( + str(error.exception), + "Unsupported layer_type: 'MonvConv'. " + "Supported types are: ARMAConv, GCSConv, GATConv, GCNConv, APPNPConv") + # Test with invalud num_layers + layer_type = "GATConv" + num_layers = 0 + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual(str( + error.exception), "num_layers must be at least 1, got 0.") + # Test with invalid num_output + num_layers = 2 + num_output = 0 + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual(str( + error.exception), "num_output must be at least 1, got 0.") + # Test with invalid num_channels + num_output = 2 + num_channels = 0 + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual(str( + error.exception), "num_channels must be at least 1, got 0.") + # Test with invalid activation + num_channels = 16 + activation = 5 + with self.assertRaises(ValueError) as error: + model = get_model(layer_type, num_layers, + num_channels, activation, num_output) + self.assertEqual( + str(error.exception), + "activation must be a string, got int.") + + def test_create_dataset(self): + + # Create dummy data in the fake filesystem for testing + num_samples = 10 + num_nodes = 5 + num_node_features = 3 + output_dim = 4 + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) + + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + self.assertEqual(len(dataset), num_samples) + for graph in dataset: + self.assertEqual(graph.x.shape, (num_nodes, num_node_features)) + self.assertEqual(dataset.a.shape, (num_nodes, num_nodes)) + self.assertEqual(graph.y.shape, (num_nodes, output_dim)) + # Clean up + self.fs.remove_object(path_cases) + self.fs.remove_object(os.path.join( + path_mobility, "commuter_mobility_2022.txt")) + self.fs.remove_object(path_mobility) + + def test_train_step(self): + + # Create a simple model for testing + model = get_model( + layer_type="GCNConv", num_layers=2, num_channels=16, + activation="relu", num_output=2) + optimizer = tf.keras.optimizers.Adam() + loss_fn = MeanAbsolutePercentageError() + + # Create dummy data + num_samples = 10 + num_nodes = 5 + num_node_features = 3 + output_dim = 2 + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + # Build the model by calling it on a batch of data + loader = MixedLoader(dataset, batch_size=4, epochs=1) + inputs, y = next(loader) + model(inputs) + + # Perform a training step + loss, acc = train_step(inputs, y, loss_fn, model, optimizer) + + # Check if the loss is a scalar tensor + self.assertIsInstance(loss, tf.Tensor) + self.assertEqual(loss.shape, ()) + self.assertIsInstance(acc, tf.Tensor) + self.assertEqual(acc.shape, ()) + self.assertGreaterEqual(acc.numpy(), 0) + self.assertGreaterEqual(loss.numpy(), 0) + + def test_evaluate(self): + + # Create a simple model for testing + model = get_model( + layer_type="GCNConv", num_layers=2, num_channels=16, + activation="relu", num_output=4) + loss_fn = MeanAbsolutePercentageError() + + # Create dummy data in the fake filesystem for testing + num_samples = 10 + num_nodes = 5 + num_node_features = 3 + output_dim = 4 + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + # Build the model by calling it on a batch of data + loader = MixedLoader(dataset, batch_size=2, epochs=1) + inputs, _ = loader.__next__() + model(inputs) + # Redefine the loader + loader = MixedLoader(dataset, batch_size=2, epochs=1) + res = evaluate(loader, model, loss_fn) + # Check if the result is a tuple of (loss, accuracy) + self.assertEqual(len(res), 2) + self.assertGreaterEqual(res[0], 0) + self.assertGreaterEqual(res[1], 0) + + # Test with retransformation + loader = MixedLoader(dataset, batch_size=2, epochs=1) + res = evaluate(loader, model, loss_fn, True) + + self.assertEqual(len(res), 2) + self.assertGreaterEqual(res[0], 0) + self.assertGreaterEqual(res[1], 0) + # Clean up + self.fs.remove_object(path_cases) + self.fs.remove_object(os.path.join( + path_mobility, "commuter_mobility_2022.txt")) + self.fs.remove_object(path_mobility) + + @patch("os.path.realpath", return_value="/home/") + def test_train_and_evaluate(self, mock_realpath): + + number_of_epochs = 2 + # Create a simple model for testing + model = get_model( + layer_type="GCNConv", num_layers=2, num_channels=16, + activation="relu", num_output=4) + + # Create dummy data in the fake filesystem for testing + num_samples = 20 + num_nodes = 5 + num_node_features = 3 + output_dim = 4 + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + + res = train_and_evaluate( + dataset, + batch_size=2, + epochs=number_of_epochs, + model=model, + loss_fn=MeanAbsolutePercentageError(), + optimizer=tf.keras.optimizers.Adam(), + es_patience=100) + + self.assertEqual(len(res["train_losses"][0]), number_of_epochs) + self.assertEqual(len(res["val_losses"][0]), number_of_epochs) + self.assertGreater(res["mean_test_loss"], 0) + + # Testing with saving the results + res = train_and_evaluate( + dataset, + batch_size=2, + epochs=number_of_epochs, + model=model, + loss_fn=MeanAbsolutePercentageError(), + optimizer=tf.keras.optimizers.Adam(), + es_patience=100, + save_dir=self.path) + save_results_path = os.path.join(self.path, "model_evaluations_paper") + save_model_path = os.path.join(self.path, "saved_weights") + self.assertTrue(os.path.exists(save_results_path)) + self.assertTrue(os.path.exists(save_model_path)) + + file_path_df = save_results_path+"/model.csv" + df = pd.read_csv(file_path_df) + self.assertEqual(len(df), 1) + for item in [ + "train_loss", "val_loss", "test_loss", + "test_loss_orig", "training_time", + "loss_history", "val_loss_history"]: + self.assertIn(item, df.columns) + + file_path_model = save_model_path+"/model.pickle" + with open(file_path_model, 'rb') as f: + weights_loaded = pickle.load(f) + weights = model.get_weights() + for w1, w2 in zip(weights_loaded, weights): + np.testing.assert_array_equal(w1, w2) + # Clean up + self.fs.remove_object(path_cases) + self.fs.remove_object(os.path.join( + path_mobility, "commuter_mobility_2022.txt")) + self.fs.remove_object(path_mobility) + self.fs.remove_object(save_results_path) + self.fs.remove_object(save_model_path) + + def test_perform_grid_search(self): + + # Create dummy data in the fake filesystem for testing + num_samples = 10 + num_nodes = 4 + num_node_features = 3 + output_dim = 4 + data = self.create_dummy_data( + num_samples, num_nodes, num_node_features, output_dim) + # Save dummy data to the fake file system + path_cases, path_mobility = self.setup_fake_filesystem( + self.fs, self.path, data) + # Create dataset + dataset = create_dataset( + path_cases, path_mobility, number_of_nodes=num_nodes) + + # Define model parameters for grid search + layers = ["GCNConv"] + num_layers = [1] + num_channels = [8] + activations = ["relu"] + parameter_grid = [(layer, n_layer, channel, activation) + for layer in layers for n_layer in num_layers + for channel in num_channels + for activation in activations] + batch_size = 2 + es_patience = 5 + max_epochs = 2 + # Perform grid search with explicit save_dir to avoid os.path.realpath issues + perform_grid_search(dataset, parameter_grid, self.path, + batch_size=batch_size, max_epochs=max_epochs, + es_patience=es_patience) + + # Check if the results file is created + results_file = os.path.join( + self.path, "saves", "grid_search_results.csv") + self.assertTrue(os.path.exists(results_file)) + + # Check if the results file has the expected number of rows + df_results = pd.read_csv(results_file) + self.assertEqual(len(df_results), len(parameter_grid)) + self.assertEqual(len(df_results.columns), 12) + + def test_scale_data_valid_data(self): + """Test utils.scale_data with valid input and label data.""" + data = { + # 10 samples, 1 day, 5 nodes, 8 groups + "inputs": np.random.rand(10, 1, 8, 5), + "labels": np.random.rand(10, 1, 8, 5) + } + + scaled_inputs, scaled_labels = utils.scale_data(data, True) + + # Check that the scaled data is not equal to the original data + assert not np.allclose( + data["inputs"].transpose(0, 3, 1, 2), scaled_inputs) + assert not np.allclose( + data["labels"].transpose(0, 3, 1, 2), scaled_labels) + + # Check that the scaled data is log-transformed + assert np.allclose(scaled_inputs, np.log1p( + data["inputs"]).transpose(0, 3, 1, 2)) + assert np.allclose(scaled_labels, np.log1p( + data["labels"]).transpose(0, 3, 1, 2)) + + def test_scale_data_invalid_data(self): + """Test utils.scale_data with invalid (non-numeric) data.""" + data = { + "inputs": np.array([["a", "b"], ["c", "d"]]), # Non-numeric data + "labels": np.array([["e", "f"], ["g", "h"]]) + } + + with self.assertRaises(ValueError): + utils.scale_data(data) + + +if __name__ == '__main__': + unittest.main()