def spread_plot(args): """Draw the spreading as a function of time.""" # Get colours and line styles from default colours = get_colours(args.colour, len(args.spreading)) labels, draw_legend = get_labels(args.label, len(args.spreading)) linestyles = {} linestyles['line'] = get_linestyles(args.linestyle, len(args.spreading)) linestyles['fit'] = get_linestyles(args.fitstyle, len(args.spreading), 'dashed') # Create linear fitting function fitfunc = lambda p, t, r: r - p[0] - p[1] * t pinit = [1.0, 1/7] # Find shift array for synchronisation shift_array = get_shift(args.spreading, sync=args.sync) impact_shift = get_shift(args.spreading, sync='impact') # Create dicts for lists of fit constants (r = amp * t**index) amp = {} index = {} ampError = {} indexError = {} ampMean = [] ampMeanError = [] indexMean = [] indexMeanError = [] for i, spread_list in enumerate(args.spreading): amp[i] = [] index[i] = [] ampError[i] = [] indexError[i] = [] spread, full_data = combine_spread(spread_list, shift=shift_array[i]) spread.times = np.array(spread.times) - spread.times[0] for k, _file in enumerate(spread_list): data = Spread().read(_file) data.times = np.array(data.times) - impact_shift[i][k] # Get radius and domain radius = {'real': np.array(calc_radius(data))} domain = {'real': np.array(data.times)} # Cut times outside of range for j, time in enumerate(domain['real']): if time > args.tend: radius['real'] = radius['real'][:j] domain['real'] = domain['real'][:j] # Add logged values radius['log'] = np.log10(radius['real'][1:]) domain['log'] = np.log10(domain['real'][1:]) # Cut in log range for j, logt in enumerate(domain['log']): if logt > args.tendlog: radius['log'] = radius['log'][:j] domain['log'] = domain['log'][:j] # Fit constants to data out = optimize.leastsq(fitfunc, pinit, args=(domain['log'], radius['log']), full_output=1) pfinal = out[0] covar = out[1] # Add unlogged constants to lists amp[i].append(10**pfinal[0]) index[i].append(pfinal[1]) ampError[i].append(np.sqrt(covar[1][1]) * amp[i][-1]) indexError[i].append(np.sqrt(covar[0][0])) if args.draw == 'log' and args.nomean: plot_line( line=radius['log'], domain=domain['log'], color=colours[i], linestyle=linestyles['line'][i] ) if not args.nofit: plot_line( line=out[0][0] + out[0][1] * domain['log'], domain=domain['log'], color=colours[i], linestyle=linestyles['fit'][i] ) if args.draw == 'real' and args.nomean: plot_line( line=radius['real'], domain=domain['real'], color=colours[i], linestyle=linestyles['line'][i] ) if not args.nofit: plot_line( line=amp[i][-1] * (domain['real']**index[i][-1]), domain=domain['real'], color=colours[i], linestyle=linestyles['fit'][i] ) ampMean.append(np.mean(amp[i])) ampMeanError.append(np.std(amp[i]) / np.sqrt(len(amp[i]) - 1)) indexMean.append(np.mean(index[i])) indexMeanError.append(np.std(index[i]) / np.sqrt(len(index[i]) - 1)) if not args.nomean: if args.draw == 'log': plot_line( line=np.log10(calc_radius(spread)), domain=np.log10(spread.times), label=labels[i], color=colours[i], linestyle=linestyles['line'][i] ) if not args.nofit: plot_line( line=(np.log10(ampMean[i]) + indexMean[i] * np.log10(spread.times)), domain=np.log10(spread.times), label='C=%.2f, n=%.2f'%(ampMean[i], indexMean[i]), color=colours[i], linestyle=linestyles['fit'][i] ) if args.draw == 'real': plot_line( line=calc_radius(spread), domain=spread.times, label=labels[i], color=colours[i], linestyle=linestyles['line'][i] ) if not args.nofit: plot_line( line=ampMean[i] * (domain['real']**indexMean[i]), domain=domain['real'], label='C=%.2f, n=%.2f'%(ampMean[i], indexMean[i]), color=colours[i], linestyle=linestyles['fit'][i] ) plt.title(args.title, fontsize='medium') plt.axis('normal') plt.legend() # Default xlabel and xlims based on draw method if args.draw == 'real': if args.ylabel == None: args.ylabel = "Spread radius (nm)" if args.xlabel == None: args.xlabel = "Time (ps)" if (args.tend and args.tendlog) < np.inf: plt.xlim([None, min(args.tend, 10**args.tendlog)]) elif args.draw == 'log': if args.ylabel == None: args.ylabel = "log10 of radius (in nm)" if args.xlabel == None: args.xlabel = "log10 of time (in ps)" if (args.tend and args.tendlog) < np.inf: plt.xlim([None, min(args.tend, args.tendlog)]) plt.xlabel(args.xlabel, fontsize='medium') plt.ylabel(args.ylabel, fontsize='medium') if args.xlim: plt.xlim(args.xlim) # Print collected output print("Fitting spread radius 'R' of input file sets to power law functions " "of time 't' as 'R = C * (t ** n)' and taking means:") for i, _ in enumerate(amp): print() # If nomean, print individual line values if args.nomean: for values in zip(amp[i], ampError[i], index[i], indexError[i]): print("%f +/- %f" % (values[0], values[1]), end=', ') print("%f +/- %f" % (values[2], values[3])) # Print mean values if args.nomean: print(" -> ", end='') print("C = %f +/- %f" % (ampMean[i], ampMeanError[i])) if args.nomean: print(" -> ", end='') print("n = %f +/- %f" % (indexMean[i], indexMeanError[i])) # Finish by saving and / or showing if args.save: plt.savefig(args.save) if args.draw != 'off': plt.show() return None
help="files to modify with new time") parser.add_argument('-dt', '--delta_t', required=True, type=float, help="new delta_t for files") parser.add_argument('-t0', type=float, help="set new initial time of maps") parser.add_argument('-o', '--output', nargs='+', help="list of files to output to instead of overwriting read") args = parser.parse_args() # Set output to output, else input if args.output: # Control that corresponding files exist if (len(args.spreading) != len(args.output)): parser.error( "number of output files does not match number of spread files" ) output = args.output else: output = args.spreading # Edit and output one at a time for i, _file in enumerate(args.spreading): spread = Spread().read(_file) spread.delta_t = args.delta_t times = np.array(spread.frames) * args.delta_t if args.t0 != None: times -= times[0] - args.t0 spread.times = list(times) spread.save(output[i])