Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add option to a minimal archive save method #16

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add option to a minimal archive save method #16

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 27 additions & 2 deletions rivus/main/rivus.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1327,7 +1327,7 @@ def save_log(result, filename):
file_handle.write(str(result))


def save(prob, filename):
def save(prob, filename, mode='all'):
"""Save rivus model instance to a gzip'ed pickle file

Pickle is the standard Python way of serializing and de-serializing Python
Expand All @@ -1343,7 +1343,13 @@ def save(prob, filename):
Args:
prob: a rivus model instance
filename: pickle file to be written

mode: what to save into the archive
all - store all the rivus model [default]
io - store inputs and outputs without the model
in a dict() with:
'costs', 'pmax', 'kappa_hub', 'kappa_process',
'source', 'flows', 'hubs', 'proc_io', 'proc_tau',
'r_out', 'r_in', 'peak'
Returns:
Nothing
"""
Expand All @@ -1353,7 +1359,26 @@ def save(prob, filename):
except ImportError:
import pickle
with gzip.GzipFile(filename, 'wb') as file_handle:
if mode == 'io':
outp = {}
for para in ('edge', 'vertex'):
outp[para] = prob.params[para]
costs, pmax, kappa_hub, kappa_process = get_constants(prob)
source, flows, hubs, proc_io, proc_tau = get_timeseries(prob)
r_out, r_in = prob.r_out, prob.r_in
peak = prob.peak
vals = [costs, pmax, kappa_hub, kappa_process,
source, flows, hubs, proc_io, proc_tau,
r_out, r_in, peak]
keys = ['costs', 'pmax', 'kappa_hub', 'kappa_process',
'source', 'flows', 'hubs', 'proc_io', 'proc_tau',
'r_out', 'r_in', 'peak']
for obj, key in zip(vals, keys):
outp[key] = obj
prob = outp

pickle.dump(prob, file_handle)


def load(filename):
"""Load a rivus model instance from a gzip'ed pickle file
Expand Down