def swap_colors(json_file_path):
'''
Switches out color ramp in meta.json files.
Uses custom color ramp if provided and valid; otherwise falls back to nextstrain default colors.
N.B.: Modifies json in place and writes to original file path.
'''
j = json.load(open(json_file_path, 'r'))
color_options = j['color_options']
for k,v in color_options.items():
if 'color_map' in v:
categories, colors = zip(*v['color_map'])
## Use custom colors if provided AND present for all categories in the dataset
if custom_colors and all([category in custom_colors for category in categories]):
colors = [ custom_colors[category] for category in categories ]
## Expand the color palette if we have too many categories
elif len(categories) > len(default_colors):
from matplotlib.colors import LinearSegmentedColormap, to_hex
from numpy import linspace
expanded_cmap = LinearSegmentedColormap.from_list('expanded_cmap', default_colors[-1], N=len(categories))
discrete_colors = [expanded_cmap(i) for i in linspace(0,1,len(categories))]
colors = [to_hex(c).upper() for c in discrete_colors]
else: ## Falls back to default nextstrain colors
colors = default_colors[len(categories)] # based on how many categories are present; keeps original ordering
j['color_options'][k]['color_map'] = map(list, zip(categories, colors))
json.dump(j, open(json_file_path, 'w'), indent=1)