Skip to content

Company saver

load_companies_from_hdf5(file_path, chunk_size=50000, domain_super_areas=None, config_filename=None)

Loads companies from an hdf5 file located at file_path. Note that this object will not be ready to use, as the links to object instances of other classes need to be restored first. This function should be rarely be called oustide world.py

Parameters:

Name Type Description Default
file_path str
required
chunk_size

(Default value = 50000)

50000
domain_super_areas

(Default value = None)

None
config_filename

(Default value = None)

None
Source code in june/hdf5_savers/company_saver.py
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def load_companies_from_hdf5(
    file_path: str, chunk_size=50000, domain_super_areas=None, config_filename=None
):
    """Loads companies from an hdf5 file located at ``file_path``.
    Note that this object will not be ready to use, as the links to
    object instances of other classes need to be restored first.
    This function should be rarely be called oustide world.py

    Args:
        file_path (str): 
        chunk_size: (Default value = 50000)
        domain_super_areas: (Default value = None)
        config_filename: (Default value = None)

    """

    Company_Class = Company
    disease_config = GlobalContext.get_disease_config()
    Company_Class.subgroup_params = SubgroupParams.from_disease_config(disease_config)

    logger.info("loading companies...")
    with h5py.File(file_path, "r", libver="latest", swmr=True) as f:
        companies = f["companies"]
        companies_list = []
        n_companies = companies.attrs["n_companies"]
        n_chunks = int(np.ceil(n_companies / chunk_size))

        # Check if subgroup registered members data exists
        has_subgroup_data = "registered_members_subgroups" in companies

        if has_subgroup_data:
            # Get all subgroups
            subgroups = read_dataset(companies["registered_members_subgroups"], 0, companies["registered_members_subgroups"].shape[0])

            # For each subgroup, prepare the count and flattened arrays
            subgroup_counts = {}
            subgroup_members = {}
            subgroup_cumulative_counts = {}

            for subgroup_id in subgroups:
                sg_key = f"registered_members_count_sg{subgroup_id}"
                ids_key = f"registered_members_ids_sg{subgroup_id}"

                if sg_key in companies and ids_key in companies:
                    # Read counts for this subgroup
                    counts = read_dataset(companies[sg_key], 0, n_companies)
                    subgroup_counts[subgroup_id] = counts

                    # Calculate cumulative counts for this subgroup
                    cumulative = np.concatenate(([0], np.cumsum(counts)))
                    subgroup_cumulative_counts[subgroup_id] = cumulative

                    # Read flattened member IDs for this subgroup
                    if companies[ids_key].shape[0] > 0:
                        member_ids = read_dataset(companies[ids_key], 0, companies[ids_key].shape[0])
                        subgroup_members[subgroup_id] = member_ids
                    else:
                        subgroup_members[subgroup_id] = np.array([])

        for chunk in range(n_chunks):
            logger.info(f"Companies chunk {chunk} of {n_chunks}")
            idx1 = chunk * chunk_size
            idx2 = min((chunk + 1) * chunk_size, n_companies)
            length = idx2 - idx1
            ids = read_dataset(companies["id"], idx1, idx2)
            sectors = read_dataset(companies["sector"], idx1, idx2)
            n_workers_maxs = read_dataset(companies["n_workers_max"], idx1, idx2)
            super_areas = read_dataset(companies["super_area"], idx1, idx2)

            for k in range(length):
                if domain_super_areas is not None:
                    super_area = super_areas[k]
                    if super_area == nan_integer:
                        raise ValueError(
                            "if ``domain_super_areas`` is True, I expect not Nones super areas."
                        )
                    if super_area not in domain_super_areas:
                        continue

                # Get registered_members_ids for this company
                registered_members_dict = {}

                if has_subgroup_data:
                    company_index = idx1 + k

                    # Get members for each subgroup
                    for subgroup_id in subgroups:
                        if subgroup_id in subgroup_counts and subgroup_id in subgroup_members:
                            n_members = subgroup_counts[subgroup_id][company_index]

                            if n_members > 0 and len(subgroup_members[subgroup_id]) > 0:
                                start_idx = subgroup_cumulative_counts[subgroup_id][company_index]
                                end_idx = subgroup_cumulative_counts[subgroup_id][company_index + 1]

                                if start_idx < len(subgroup_members[subgroup_id]):
                                    registered_members_dict[int(subgroup_id)] = subgroup_members[subgroup_id][start_idx:end_idx].tolist()

                company = Company(
                    super_area=None,
                    n_workers_max=n_workers_maxs[k],
                    sector=sectors[k].decode(),
                    registered_members_ids=registered_members_dict,
                )
                company.id = ids[k]
                companies_list.append(company)
    return Companies(companies_list)

restore_companies_properties_from_hdf5(world, file_path, chunk_size, domain_super_areas=None)

Parameters:

Name Type Description Default
world World
required
file_path str
required
chunk_size
required
domain_super_areas

(Default value = None)

None
Source code in june/hdf5_savers/company_saver.py
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
def restore_companies_properties_from_hdf5(
    world: World, file_path: str, chunk_size, domain_super_areas=None
):
    """

    Args:
        world (World): 
        file_path (str): 
        chunk_size: 
        domain_super_areas: (Default value = None)

    """
    with h5py.File(file_path, "r", libver="latest", swmr=True) as f:
        companies = f["companies"]
        n_companies = companies.attrs["n_companies"]
        n_chunks = int(np.ceil(n_companies / chunk_size))

        # Check if subgroup registered members data exists
        has_subgroup_data = "registered_members_subgroups" in companies

        if has_subgroup_data:
            # Get all subgroups
            subgroups = read_dataset(companies["registered_members_subgroups"], 0, companies["registered_members_subgroups"].shape[0])

            # For each subgroup, prepare the count and flattened arrays
            subgroup_counts = {}
            subgroup_members = {}
            subgroup_cumulative_counts = {}

            for subgroup_id in subgroups:
                sg_key = f"registered_members_count_sg{subgroup_id}"
                ids_key = f"registered_members_ids_sg{subgroup_id}"

                if sg_key in companies and ids_key in companies:
                    # Read counts for this subgroup
                    counts = read_dataset(companies[sg_key], 0, n_companies)
                    subgroup_counts[subgroup_id] = counts

                    # Calculate cumulative counts for this subgroup
                    cumulative = np.concatenate(([0], np.cumsum(counts)))
                    subgroup_cumulative_counts[subgroup_id] = cumulative

                    # Read flattened member IDs for this subgroup
                    if companies[ids_key].shape[0] > 0:
                        member_ids = read_dataset(companies[ids_key], 0, companies[ids_key].shape[0])
                        subgroup_members[subgroup_id] = member_ids
                    else:
                        subgroup_members[subgroup_id] = np.array([])

        for chunk in range(n_chunks):
            idx1 = chunk * chunk_size
            idx2 = min((chunk + 1) * chunk_size, n_companies)
            length = idx2 - idx1
            ids = read_dataset(companies["id"], idx1, idx2)
            super_areas = read_dataset(companies["super_area"], idx1, idx2)

            for k in range(length):
                if domain_super_areas is not None:
                    super_area = super_areas[k]
                    if super_area == nan_integer:
                        raise ValueError(
                            "if ``domain_super_areas`` is True, I expect not Nones super areas."
                        )
                    if super_area not in domain_super_areas:
                        continue

                company = world.companies.get_from_id(ids[k])

                if super_areas[k] == nan_integer:
                    company.super_area = None
                else:
                    company.super_area = world.super_areas.get_from_id(super_areas[k])
                # Restore registered_members_ids if available
                if has_subgroup_data:
                    company_index = idx1 + k
                    registered_members_dict = {}

                    # Process each subgroup
                    for subgroup_id in subgroups:
                        if subgroup_id in subgroup_counts and subgroup_id in subgroup_members:
                            n_members = subgroup_counts[subgroup_id][company_index]

                            if n_members > 0 and len(subgroup_members[subgroup_id]) > 0:
                                start_idx = subgroup_cumulative_counts[subgroup_id][company_index]
                                end_idx = subgroup_cumulative_counts[subgroup_id][company_index + 1]

                                if start_idx < len(subgroup_members[subgroup_id]):
                                    registered_members_dict[int(subgroup_id)] = subgroup_members[subgroup_id][start_idx:end_idx].tolist()

                    # Always use dictionary format
                    company.registered_members_ids = registered_members_dict

save_companies_to_hdf5(companies, file_path, chunk_size=500000)

Saves the Population object to hdf5 format file file_path. Currently for each person, the following values are stored: - id, super_area, sector, n_workers_max,

Parameters:

Name Type Description Default
companies Companies

population object

required
file_path str

path of the saved hdf5 file

required
chunk_size int

number of people to save at a time. Note that they have to be copied to be saved,

500000

so keep the number below 1e6. (Default value = 500000)

Source code in june/hdf5_savers/company_saver.py
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def save_companies_to_hdf5(
    companies: Companies, file_path: str, chunk_size: int = 500000
):
    """Saves the Population object to hdf5 format file ``file_path``. Currently for each person,
    the following values are stored:
    - id, super_area, sector, n_workers_max,

    Args:
        companies (Companies): population object
        file_path (str): path of the saved hdf5 file
        chunk_size (int, optional): number of people to save at a time. Note that they have to be copied to be saved,
    so keep the number below 1e6. (Default value = 500000)

    """
    n_companies = len(companies)
    n_chunks = int(np.ceil(n_companies / chunk_size))
    with h5py.File(file_path, "a") as f:
        companies_dset = f.create_group("companies")
        first_company_idx = companies[0].id
        for chunk in range(n_chunks):
            idx1 = chunk * chunk_size
            idx2 = min((chunk + 1) * chunk_size, n_companies)
            ids = []
            super_areas = []
            sectors = []
            n_workers_max = []
            registered_members_ids = []  # Add new list for registered_members_ids
            company_idx = [company.id for company in companies[idx1:idx2]]
            # sort companies by id
            companies_sorted = [
                companies[i - first_company_idx] for i in np.sort(company_idx)
            ]
            for company in companies_sorted:
                ids.append(company.id)
                if company.super_area is None:
                    super_areas.append(nan_integer)
                else:
                    super_areas.append(company.super_area.id)
                sectors.append(company.sector.encode("ascii", "ignore"))
                n_workers_max.append(company.n_workers_max)
                # Handle registered_members_ids - always a dictionary with subgroups
                if hasattr(company, 'registered_members_ids') and company.registered_members_ids is not None:
                    subgroup_dict = {}
                    for subgroup_id, members in company.registered_members_ids.items():
                        subgroup_dict[subgroup_id] = np.array(members, dtype=np.int64)
                    registered_members_ids.append(subgroup_dict)
                else:
                    registered_members_ids.append({})  # Empty dictionary

            ids = np.array(ids, dtype=np.int64)
            super_areas = np.array(super_areas, dtype=np.int64)
            sectors = np.array(sectors, dtype="S10")
            n_workers_max = np.array(n_workers_max, dtype=np.float64)

            # Create datasets to store the number of registered members per company per subgroup
            # First, determine what subgroups exist across all companies
            all_subgroups = set()
            for company_subgroups in registered_members_ids:
                all_subgroups.update(company_subgroups.keys())
            all_subgroups = sorted(list(all_subgroups))  # Ensure consistent ordering

            # Store subgroup IDs as integers
            subgroup_ids = np.array(all_subgroups, dtype=np.int64)

            # Create counts and flattened arrays for each subgroup
            subgroup_counts = {}
            flattened_subgroup_members = {}

            for subgroup_id in all_subgroups:
                # Count of members in this subgroup for each company
                counts = np.array([len(company_dict.get(subgroup_id, [])) for company_dict in registered_members_ids], dtype=np.int64)
                subgroup_counts[subgroup_id] = counts

                # Flatten all member IDs for this subgroup
                member_arrays = [company_dict.get(subgroup_id, np.array([], dtype=np.int64)) for company_dict in registered_members_ids]
                flattened = np.concatenate(member_arrays) if any(len(arr) > 0 for arr in member_arrays) else np.array([], dtype=np.int64)
                flattened_subgroup_members[subgroup_id] = flattened
            if chunk == 0:
                companies_dset.attrs["n_companies"] = n_companies
                companies_dset.create_dataset("id", data=ids, maxshape=(None,))
                companies_dset.create_dataset(
                    "super_area", data=super_areas, maxshape=(None,)
                )
                companies_dset.create_dataset("sector", data=sectors, maxshape=(None,))
                companies_dset.create_dataset(
                    "n_workers_max", data=n_workers_max, maxshape=(None,)
                )
                # Store registered_members_ids subgroups information
                # Store the subgroup IDs
                companies_dset.create_dataset("registered_members_subgroups", data=subgroup_ids, maxshape=(None,))

                # Create datasets for each subgroup
                for subgroup_id in all_subgroups:
                    # Store counts for this subgroup
                    companies_dset.create_dataset(
                        f"registered_members_count_sg{subgroup_id}", 
                        data=subgroup_counts[subgroup_id], 
                        maxshape=(None,)
                    )

                    # Store flattened members for this subgroup
                    flattened_members = flattened_subgroup_members[subgroup_id]
                    companies_dset.create_dataset(
                        f"registered_members_ids_sg{subgroup_id}", 
                        data=flattened_members, 
                        maxshape=(None,)
                    )
            else:
                newshape = (companies_dset["id"].shape[0] + ids.shape[0],)
                companies_dset["id"].resize(newshape)
                companies_dset["id"][idx1:idx2] = ids
                companies_dset["super_area"].resize(newshape)
                companies_dset["super_area"][idx1:idx2] = super_areas
                companies_dset["sector"].resize(newshape)
                companies_dset["sector"][idx1:idx2] = sectors
                companies_dset["n_workers_max"].resize(newshape)
                companies_dset["n_workers_max"][idx1:idx2] = n_workers_max

                # Update registered members for subgroups
                for subgroup_id in all_subgroups:
                    # Update counts for this subgroup
                    count_dataset = companies_dset[f"registered_members_count_sg{subgroup_id}"]
                    count_dataset.resize(newshape)
                    count_dataset[idx1:idx2] = subgroup_counts[subgroup_id]

                    # Update flattened IDs for this subgroup (variable length)
                    ids_dataset = companies_dset[f"registered_members_ids_sg{subgroup_id}"]
                    flattened_members = flattened_subgroup_members[subgroup_id]

                    if flattened_members.shape[0] > 0:
                        current_length = ids_dataset.shape[0]
                        new_length = current_length + flattened_members.shape[0]
                        ids_dataset.resize((new_length,))
                        ids_dataset[current_length:new_length] = flattened_members