diff --git a/gwemopt/args.py b/gwemopt/args.py index 80516e0..8f547da 100644 --- a/gwemopt/args.py +++ b/gwemopt/args.py @@ -152,4 +152,8 @@ def parse_args(args): parser.add_argument("--inclination", action="store_true", default=False) parser.add_argument("--projection", default="astro mollweide") + parser.add_argument("--solverType", default="heuristic") + parser.add_argument("--milpSolver", default="PULP_CBC_CMD") + parser.add_argument("--milpOptions", default="{}") + return parser.parse_args(args=args) diff --git a/gwemopt/params.py b/gwemopt/params.py index 8cf3386..07f04be 100644 --- a/gwemopt/params.py +++ b/gwemopt/params.py @@ -1,3 +1,4 @@ +import json import os from pathlib import Path @@ -151,4 +152,14 @@ def params_struct(opts): opts.projection if hasattr(opts, "projection") else "astro mollweide" ) + params["solverType"] = ( + opts.solverType if hasattr(opts, "solverType") else "heuristic" + ) + params["milpSolver"] = ( + opts.milpSolver if hasattr(opts, "milpSolver") else "PULP_CBC_CMD" + ) + params["milpOptions"] = ( + json.loads(opts.milpOptions) if hasattr(opts, "milpOptions") else {} + ) + return params diff --git a/gwemopt/scheduler.py b/gwemopt/scheduler.py index 0bf3e9c..b60ee2e 100644 --- a/gwemopt/scheduler.py +++ b/gwemopt/scheduler.py @@ -6,10 +6,10 @@ import ligo.segments as segments import numpy as np from astropy.time import Time -from munkres import Munkres, make_cost_matrix +from ortools.linear_solver import pywraplp from gwemopt.tiles import balance_tiles, optimize_max_tiles, schedule_alternating -from gwemopt.utils import angular_distance +from gwemopt.utils import angular_distance, solve_milp def get_altaz_tile(ra, dec, observer, obstime): @@ -88,7 +88,7 @@ def find_tile( return idx2, exposureids, probs -def get_order( +def get_order_heuristic( params, tile_struct, tilesegmentlists, exposurelist, observer, config_struct ): """ @@ -393,18 +393,62 @@ def get_order( tilematrix_mask = tilematrix > 10 ** (-10) if tilematrix_mask.any(): - print("Calculating Hungarian solution...") + print("Calculating MILP solution...") total_cost = 0 - cost_matrix = make_cost_matrix(tilematrix) - m = Munkres() - optimal_points = m.compute(cost_matrix) - print("Hungarian solution calculated...") + + maximum = max(max(row) for row in tilematrix) + inversion_function = lambda x: maximum - x + + cost_matrix = [] + for row in tilematrix: + cost_matrix.append([inversion_function(value) for value in row]) + + # Create a linear solver + solver = pywraplp.Solver.CreateSolver("CBC") + + # Define variables + num_workers = len(cost_matrix) + num_tasks = len(cost_matrix[0]) + x = {} + + for i in range(num_workers): + for j in range(num_tasks): + x[i, j] = solver.BoolVar("x[%i,%i]" % (i, j)) + + # Define constraints: each worker is assigned to at most one task + for i in range(num_workers): + solver.Add(solver.Sum([x[i, j] for j in range(num_tasks)]) <= 1) + + # Define constraints: each task is assigned to exactly one worker + for j in range(num_tasks): + solver.Add(solver.Sum([x[i, j] for i in range(num_workers)]) == 1) + + # Define objective function: minimize the total cost + objective = solver.Objective() + for i in range(num_workers): + for j in range(num_tasks): + objective.SetCoefficient(x[i, j], cost_matrix[i][j]) + objective.SetMinimization() + + # Solve the problem + status = solver.Solve() + + optimal_points = [] + if status == pywraplp.Solver.OPTIMAL: + print("Total cost =", solver.Objective().Value()) + for i in range(num_workers): + for j in range(num_tasks): + if x[i, j].solution_value() > 0: + optimal_points.append((i, j)) + else: + print("The problem does not have an optimal solution.") + max_no_observ = min(tilematrix.shape) for jj in range(max_no_observ): - idx0, idx1 = optimal_points[jj] # idx0 indexes over the time windows, idx1 indexes over the probabilities # idx2 gets the exposure id of the tile, used to assign tileexptime and tilenexps try: + idx0, idx1 = optimal_points[jj] idx2 = exposureids[idx1] pamw = tilematrix[idx0][idx1] total_cost += pamw @@ -427,6 +471,127 @@ def get_order( return idxs, filts +def get_order_milp(params, tile_struct, exposurelist, observer, config_struct): + """ + tile_struct: dictionary. key -> struct info. + exposurelist: list of segments that the telescope is supposed to be working. + consecutive segments from the start to the end, with each segment size + being the exposure time. + Returns a list of tile indices in the order of observation. + """ + + if "dec_constraint" in config_struct: + dec_constraint = config_struct["dec_constraint"].split(",") + dec_min = float(dec_constraint[0]) + dec_max = float(dec_constraint[1]) + + exposureids = [] + probs = [] + ras, decs, filts, keys = [], [], [], [] + for ii, key in enumerate(list(tile_struct.keys())): + if tile_struct[key]["prob"] == 0: + continue + if "dec_constraint" in config_struct: + if (tile_struct[key]["dec"] < dec_min) or ( + tile_struct[key]["dec"] > dec_max + ): + continue + + exposureids.append(key) + probs.append(tile_struct[key]["prob"]) + ras.append(tile_struct[key]["ra"]) + decs.append(tile_struct[key]["dec"]) + filts.append(tile_struct[key]["filt"]) + keys.append(key) + + fields = -1 * np.ones( + (len(exposurelist)), + ) + filters = ["n"] * len(exposurelist) + + if len(probs) == 0: + return fields, filters + + probs = np.array(probs) + ras = np.array(ras) + decs = np.array(decs) + exposureids = np.array(exposureids) + keys = np.array(keys) + tilematrix = np.zeros((len(exposurelist), len(ras))) + for ii in np.arange(len(exposurelist)): + # first, create an array of airmass-weighted probabilities + t = Time(exposurelist[ii][0], format="mjd") + altazs = [get_altaz_tile(ra, dec, observer, t) for ra, dec in zip(ras, decs)] + alts = np.array([altaz[0] for altaz in altazs]) + horizon = config_struct["horizon"] + horizon_mask = alts <= horizon + airmass = 1 / np.cos((90.0 - alts) * np.pi / 180.0) + airmass_mask = airmass > params["airmass"] + + airmass_weight = 10 ** (0.4 * 0.1 * (airmass - 1)) + + if params["scheduleType"] in ["greedy", "greedy_slew"]: + tilematrix[ii, :] = np.array(probs) + elif params["scheduleType"] == ["airmass_weighted", "airmass_weighted_slew"]: + tilematrix[ii, :] = np.array(probs / airmass_weight) + tilematrix[ii, horizon_mask] = np.nan + tilematrix[ii, airmass_mask] = np.nan + + for jj, key in enumerate(keys): + tilesegmentlist = tile_struct[key]["segmentlist"] + if not tilesegmentlist.intersects_segment(exposurelist[ii]): + tilematrix[ii, jj] = np.nan + + # which fields are never observable + ind = np.where(np.nansum(tilematrix, axis=0) > 0)[0] + + probs = np.array(probs)[ind] + ras = np.array(ras)[ind] + decs = np.array(decs)[ind] + exposureids = np.array(exposureids)[ind] + filts = [filts[i] for i in ind] + tilematrix = tilematrix[:, ind] + + # which times do not have any observability + ind = np.where(np.nansum(tilematrix, axis=1) > 0)[0] + tilematrix = tilematrix[ind, :] + + cost_matrix = tilematrix + cost_matrix[np.isnan(cost_matrix)] = -np.inf + + distmatrix = np.zeros((len(ras), len(ras))) + for ii, (r, d) in enumerate(zip(ras, decs)): + dist = angular_distance(r, d, ras, decs) + if "slew" in params["scheduleType"]: + dist = dist / config_struct["slew_rate"] + dist = dist - config_struct["readout"] + dist[dist < 0] = 0 + else: + distmatrix[ii, :] = dist + distmatrix = distmatrix / np.max(distmatrix) + + dt = int(np.ceil((exposurelist[1][0] - exposurelist[0][0]) * 86400)) + optimal_points = solve_milp( + cost_matrix, + dist_matrix=distmatrix, + useDistance=False, + max_tasks_per_worker=len(params["filters"]), + useTaskSepration=False, + min_task_separation=int(np.ceil(dt / params["mindiff"])), + ) + + for optimal_point in optimal_points: + idx = ind[optimal_point[0]] + idy = optimal_point[1] + if len(filts[idy]) > 0: + fields[idx] = exposureids[idy] + filters[idx] = filts[idy][0] + filt = filts[idy][1:] + filts[idy] = filt + + return fields, filters + + def scheduler(params, config_struct, tile_struct): """ config_struct: the telescope configurations @@ -458,9 +623,17 @@ def scheduler(params, config_struct, tile_struct): for key in keys: # segments.py: tile_struct[key]["segmentlist"] is a list of segments when the tile is available for observation tilesegmentlists.append(tile_struct[key]["segmentlist"]) - keys, filts = get_order( - params, tile_struct, tilesegmentlists, exposurelist, observer, config_struct - ) + + if params["solverType"] == "heuristic": + keys, filts = get_order_heuristic( + params, tile_struct, tilesegmentlists, exposurelist, observer, config_struct + ) + elif params["solverType"] == "milp": + keys, filts = get_order_milp( + params, tile_struct, exposurelist, observer, config_struct + ) + else: + raise ValueError(f'Unknown solverType {params["solverType"]}') if params["doPlots"]: from gwemopt.plotting import make_schedule_plots diff --git a/gwemopt/utils/__init__.py b/gwemopt/utils/__init__.py index 9fbf2df..033e505 100644 --- a/gwemopt/utils/__init__.py +++ b/gwemopt/utils/__init__.py @@ -1,4 +1,5 @@ from gwemopt.utils.geometry import angular_distance +from gwemopt.utils.milp import solve_milp from gwemopt.utils.misc import auto_rasplit, get_exposures, integrationTime from gwemopt.utils.observability import calculate_observability from gwemopt.utils.param_utils import readParamsFromFile diff --git a/gwemopt/utils/milp.py b/gwemopt/utils/milp.py new file mode 100644 index 0000000..7ee0022 --- /dev/null +++ b/gwemopt/utils/milp.py @@ -0,0 +1,125 @@ +import time + +import numpy as np +import pulp +from tqdm import tqdm + + +def solve_milp( + cost_matrix, + max_tasks_per_worker=1, + useTaskSepration=False, + min_task_separation=1, + useDistance=False, + dist_matrix=None, + timeLimit=300, + milpSolver="PULP_CBC_CMD", + milpOptions={}, +): + + cost_matrix_mask = cost_matrix > 10 ** (-10) + optimal_points = [] + + if cost_matrix_mask.any(): + print("Calculating MILP solution...") + + # Create a CP-SAT model + problem = pulp.LpProblem("problem", pulp.LpMaximize) + + # Define variables + num_exposures, num_fields = cost_matrix.shape + + print("Define decision variables...") + # Define decision variables + x = { + (i, j): pulp.LpVariable(f"x_{i}_{j}", cat=pulp.LpBinary) + for i in tqdm(range(num_exposures)) + for j in range(num_fields) + } + + print("Define binary variables for task separation violation...") + # Define binary variables for task separation violation + s = { + (i, j): pulp.LpVariable(f"s_{i}_{j}", cat=pulp.LpBinary) + for i in tqdm(range(num_exposures)) + for j in range(num_fields) + } + + obj = pulp.lpSum( + x[i, j] * cost_matrix[i][j] + for i in range(num_exposures) + for j in range(num_fields) + ) + + if useDistance: + products = [ + (x[i, j], dist_matrix[j, k]) + for i in range(num_exposures) + for j in range(num_fields) + for k in range(dist_matrix.shape[1]) + ] + total_distance = pulp.LpAffineExpression(products) + obj -= total_distance + + # One field per exposure + for i in range(num_exposures): + problem += pulp.lpSum(x[i, j] for j in range(num_fields)) == 1 + + # Limit the number of tasks each worker can handle (if applicable) + for j in range(num_fields): + problem += ( + pulp.lpSum(x[i, j] for i in range(num_exposures)) + <= max_tasks_per_worker + ) + + print("Add constraints to exclude impossible assignments...") + # Add constraints to exclude impossible assignments + for i in tqdm(range(num_exposures)): + for j in range(num_fields): + if not np.isfinite(cost_matrix[i][j]): + problem += x[i, j] == 0 + + if useTaskSepration: + print("Define constraints: enforce minimum task separation...") + ## Define constraints: enforce minimum task separation + for i in tqdm(range(num_exposures)): + for j in range(num_fields): + for k in range(j + 1, num_fields): + problem += s[i, j] >= x[i, k] - x[i, j] - min_task_separation + problem += s[i, j] >= x[i, j] - x[i, k] - min_task_separation + + print("Number of variables:", len(problem.variables())) + print("Number of constraints:", len(problem.constraints)) + + time_limit = 60 # Stop the solver after 60 seconds + + if milpSolver == "PULP_CBC_CMD": + solver = pulp.getSolver(milpSolver) + elif milpSolver == "GUROBI": + solver = pulp.getSolver( + "GUROBI", manageEnv=True, envOptions=milpOptions, mip_gap=1e-6 + ) + else: + raise ValueError("milpSolver must be either PULP_CBC_CMD or GUROBI") + + solver.timeLimit = timeLimit + # solver.msg = True + status = problem.solve(solver) + + optimal_points = [] + if status in [pulp.LpStatusOptimal]: + for i in range(num_exposures): + for j in range(num_fields): + if ( + pulp.value(x[i, j]) is not None + and pulp.value(x[i, j]) > 0.5 + and np.isfinite(cost_matrix[i][j]) + ): + optimal_points.append((i, j)) + else: + print("The problem does not have a solution.") + + else: + print("The localization is not visible from the site.") + + return optimal_points diff --git a/pyproject.toml b/pyproject.toml index 5f353b1..aee2e19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,13 +44,15 @@ dependencies = [ 'joblib', 'ligo-segments', 'ligo.skymap', - "lxml", + 'lxml', 'h5py', 'healpy', + 'ortools', 'pandas', + 'pulp', 'shapely', - "tables", - "regions" + 'tables', + 'regions' ] dynamic = ["version"] @@ -114,7 +116,7 @@ parentdir_prefix = "gwemopt-" addopts = "--verbose -r s" [tool.bumpver] -current_version = "0.2.1" +current_version = "0.2.2" version_pattern = "MAJOR.MINOR.PATCH[PYTAGNUM]" commit_message = "bump version {old_version} -> {new_version}" commit = true