text:
#%%
"""
    abm_scheduling is an OpenSource python package for the analysis of
    the Nurse Scheduling Problem using an Agent Based Modeling approach
    based on
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published
    by the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.
"""
#%config IPCompleter.greedy=True

import networkx as nx
import numpy as np
import numpy.random as rnd
import collections
import math
from prettytable import PrettyTable
from random import shuffle
import copy

import matplotlib.pylab as plt
from scipy.integrate import simps
#%matplotlib inline
#plt.style.use('ggplot')

#import warnings
#warnings.filterwarnings('ignore')

#import matplotlib.mlab as mlab
import time
from datetime import datetime

#%%
class Shift():
    def __init__(self, day, shift_num, num_nurses_needed):
        self.day = day
        self.shift_num = shift_num
        self.num_nurses_needed = num_nurses_needed
        self.nurses = []

    def get_id_tuple(self):
        return (self.day, self.shift_num)

    def get_list_of_nurses_names(self):
        return [n.id_name for n in self.nurses]

class Utility_Function_Parameters():
    def __init__(self):
        self.utility_function = 'default'
        self.penalty_overbooking = 1000
        self.penalty_multiple_shifts_one_day = 1000
        self.penalty_max_continuous_days_working = 1000

#%%
class Schedule():
    num_shifts_per_day = 3
    days = ['Mo', 'Tu', 'We', 'Th', 'Fr', 'Sa', 'So']

    def __init__(self, num_nurses_needed = 1, matrix_nurses_needed = [], is_random=False):
        self.num_nurses_needed = num_nurses_needed
        self.is_random = is_random
        self.schedule = []
        i = 0
        for day in self.days:
            for shift_num in range(self.num_shifts_per_day):
                if is_random:
                    num_needed = int(rnd.normal(self.num_nurses_needed + 1))
                elif len(matrix_nurses_needed) > 0:
                    num_needed = matrix_nurses_needed[i]
                    i += 1
                else:
                    num_needed = self.num_nurses_needed

                self.schedule.append(Shift(day, shift_num + 1, num_needed))

    def get_shift_coverage(self):
        total_shifts = 0
        filled_shifts = 0
        for shift in self.schedule:
            total_shifts += shift.num_nurses_needed
            filled_shifts += min(len(shift.nurses), shift.num_nurses_needed)
        return filled_shifts / total_shifts

    def add_nurse_to_shift(self, nurse, shift_preference, add_to_nurse):
        idx_in_schedule = self.days.index(shift_preference[0]) * 3 + (shift_preference[1] - 1)
        self.schedule[idx_in_schedule].nurses.append(nurse)
        if add_to_nurse: nurse.shifts.append(shift_preference)

    def remove_nurse_from_shift(self, nurse, shift_preference, remove_from_nurse):
        idx_in_schedule = self.days.index(shift_preference[0]) * 3 + (shift_preference[1] - 1)
        for idx, n in enumerate(self.schedule[idx_in_schedule].nurses):
            if n.id_name == nurse.id_name:
                self.schedule[idx_in_schedule].nurses.pop(idx)
        if remove_from_nurse: nurse.shifts.remove(shift_preference)

    def print_filled_in_schedule(self, schedule_strs, title=''):
        t = PrettyTable([''] + self.days)
        t.align = 'l'
        for i in range(self.num_shifts_per_day):
            t.add_row([f'shift {i + 1}'] + schedule_strs[i::self.num_shifts_per_day])
        print(title)
        print(t)

    def print_schedule(self, schedule_name = ''):
        schedule_strs = []
        for shift in self.schedule:
            nurses_str = ''
            for idx, n in enumerate(shift.get_list_of_nurses_names()):
                if idx % 4 == 0:
                    nurses_str += '\n'
                nurses_str += str(n)
                if idx != len(shift.get_list_of_nurses_names()) - 1:
                    nurses_str += ','
            schedule_strs.append(
                f"need: {shift.num_nurses_needed}\n"
                f"nurses: {nurses_str}"
            )
        title = "Week's Schedule " + schedule_name
        self.print_filled_in_schedule(schedule_strs, title=title)

    def print_shift_coverage(self, schedule_name = ''):
        schedule_strs = []
        for shift in self.schedule:
            #filled_shifts += min(len(shift.nurses), shift.num_nurses_needed)
            schedule_strs.append(
                f"need: {shift.num_nurses_needed}\n"
                f"nurses: {len(shift.nurses)}\n"
                f"({(len(shift.nurses) / max(shift.num_nurses_needed,1)):.3g})"
            )
        title = "Shift Coverage " + schedule_name
        self.print_filled_in_schedule(schedule_strs, title=title)

#%%
class Nurse():
    max_working_days = 6

    def __init__(self, id_name):
        """
        Generated with default Agent Satisfaction Parameters
        """
        self.id_name = id_name
        self.shift_preferences = []
        self.shifts = []
        self.degree_of_availability = 0
        self.minimum_shifts = 0
        self.maximum_shifts = 0
        # Agent_Satisfaction_Parameters()
        self.satisfaction_function = 'default'
        self.base_decrease_coverage_satisfaction = 0.8
        self.sensibility_to_increase_assignment = 7
        self.gain_over_increase_assignment = 0
        self.value_under_assignment = 1000
        self.value_over_assignment = 1000
        self.penalty_different_shifts_in_week = 700
        self.penalty_wrong_shift_assigment = 700

    def generate_shift_preferences(self, degree_of_agent_availability, works_weekends):
        schedule = Schedule(0)
        for shift in schedule.schedule:
            if not works_weekends and shift.day in schedule.days[-2:]:
                continue
            if rnd.uniform() < degree_of_agent_availability:
                self.shift_preferences.append((shift.day, shift.shift_num))
        self.degree_of_availability = degree_of_agent_availability
        self.minimum_shifts = min(round(self.degree_of_availability*21,0), self.max_working_days)
        self.maximum_shifts = min(round(self.degree_of_availability*21,0), self.max_working_days)
        self.gain_over_increase_assignment = 1000/self.minimum_shifts
        self.sensibility_to_increase_assignment = 7 - self.maximum_shifts

    def assign_shift_preferences(self, matrix_nurse_availability = [], minimum_shifts = 0, maximum_shifts = 0):
        schedule = Schedule(0)
        i = -1
        for shift in schedule.schedule:
            i += 1
            if matrix_nurse_availability[i] == 'x':
                self.shift_preferences.append((shift.day, shift.shift_num))
                self.degree_of_availability += 1/21
        self.minimum_shifts = min(round(self.degree_of_availability*21,0), self.max_working_days) if minimum_shifts == 0 else minimum_shifts
        self.maximum_shifts = min(round(self.degree_of_availability*21,0), self.max_working_days) if maximum_shifts == 0 else maximum_shifts
        self.gain_over_increase_assignment = 1000/(np.mean([self.minimum_shifts,self.maximum_shifts]))
        self.sensibility_to_increase_assignment = 7 - self.maximum_shifts

    def print_shift_preferences(self):
        schedule = Schedule(0)
        schedule_strs = []
        for shift in schedule.schedule:
            if shift.get_id_tuple() in self.shift_preferences:
                schedule_strs.append('x')
            else:
                schedule_strs.append(' ')
        title = f"Nurse {self.id_name}'s Preferences. Availability: " + str(len(self.shift_preferences)) + f"({self.degree_of_availability:.3g})" +  " Min/Max: " + str(self.minimum_shifts) +" / " +str(self.maximum_shifts)
        schedule.print_filled_in_schedule(schedule_strs, title=title)

    def print_assigned_shifts(self):
        schedule = Schedule(0)
        schedule_strs = []
        for shift in schedule.schedule:
            if shift.get_id_tuple() in self.shifts:
                schedule_strs.append('1')
            else:
                schedule_strs.append(' ')
        schedule.print_filled_in_schedule(schedule_strs, title=f"Nurse {self.id_name}'s Assigned Shifts")

    def get_satisfaction(self):
        if self.satisfaction_function == 'default':
            return self.get_satisfaction_default_function()
        return 1

    def get_satisfaction_default_function(self):
        """
        Default function for agent satisfaction that accounts for:
        - coverage of minimum and maximum desired shift assignments,
        - satisfaction for new shift assignment decreases incrementally as coverages approaches the maximum
        """
        # formula for satisfaction from shift assignment =
        # base_decrease_coverage_satisfaction ^ (number_shifts - sensibility_to_increase_assignment)
        # x = assigned shifts
        s = lambda x: pow(self.base_decrease_coverage_satisfaction, (x - self.sensibility_to_increase_assignment))

        nr_shifts = len(self.shifts)
        cummulated_satisfaction_from_assigned_shifts = 0
        rate_under_assignment = (-1 + nr_shifts/self.minimum_shifts)
        rate_over_assignment = (self.maximum_shifts - nr_shifts)/self.maximum_shifts
        # matching_assigned_shifts = len(set(self.shifts).intersection(self.shift_preferences))
        # penalty when under assignment
        cummulated_satisfaction_from_assigned_shifts += self.value_under_assignment * min(0, rate_under_assignment)
        # penalty when over assignment
        cummulated_satisfaction_from_assigned_shifts += self.value_over_assignment * min(0, rate_over_assignment)
        # increase satisfaction on number of matching assigned shifts
        # we asume all assigned shifts here match because they are choosen only from their preferences
        if nr_shifts <=self.maximum_shifts:
            x_vals = np.array(list(range(nr_shifts)))
            cummulated_satisfaction_from_assigned_shifts += self.gain_over_increase_assignment * sum(s(x_vals))

        return cummulated_satisfaction_from_assigned_shifts

class NSP_AB_Model_Run_Results():
    p_to_accept_negative_change = 0
    best_schedule = []
    utility_each_timestep = []
    shift_coverage_each_timestep = []
    shift_coverage = 0
    utility = 0
    utility_function = ''
    total_agent_satisfaction = 0
    nurses = []

#%%
class NSP_AB_Model():
    def generate_nurses(self, num_nurses, degree_of_agent_availability, works_weekends):
        nurses = []
        for n in range(num_nurses):
            nurse = Nurse(id_name=n)
            nurse.generate_shift_preferences(degree_of_agent_availability, works_weekends)
            nurses.append(nurse)
        return nurses

    def generate_nurses_from_nurse_schedules(self, list_of_nurse_schedules):
        nurses = []
        i = -1
        for s in list_of_nurse_schedules:
            i += 1
            nurse = Nurse(id_name=i)
            nurse.assign_shift_preferences(s)
            nurses.append(nurse)
        return nurses

    def show_hypothetical_max_schedule(self, schedule, nurses, print_detail_schedule = True, print_shift_coverage = True, print_nurse_productivity = False):
        # show what all shift preferences look like in the schedule
        hypothetical_max_schedule = copy.deepcopy(schedule)
        hiypotetical_max_nurses = copy.deepcopy(nurses)
        for nurse in hiypotetical_max_nurses:
            for shift in nurse.shift_preferences:
                hypothetical_max_schedule.add_nurse_to_shift(nurse, shift, True)

        print('Crude hypothetical shift coverage:', hypothetical_max_schedule.get_shift_coverage())
        if print_detail_schedule:
            hypothetical_max_schedule.print_schedule(schedule_name="Hypothetical Maximum")
        if print_shift_coverage:
            hypothetical_max_schedule.print_shift_coverage(schedule_name="Hypothetical Maximum")
        if print_nurse_productivity:
            self.print_nurse_productivity(nurses = hiypotetical_max_nurses)

    def get_utility(self, schedule: Schedule, nurses: [Nurse], utility_function_parameters: Utility_Function_Parameters):
        if utility_function_parameters.utility_function == 'default':
            return self.get_utility_default(schedule = schedule, utility_function_parameters=utility_function_parameters)
        if utility_function_parameters.utility_function == 'agent_satisfaction':
            return self.get_utility_agent_satisfaction(schedule = schedule, nurses=nurses, utility_function_parameters=utility_function_parameters)

    def get_utility_default(self, schedule: Schedule, utility_function_parameters: Utility_Function_Parameters):
        """
        Default utility function for the Shift assignment considering:
        - negative penalty for overbooking on the shifts
        - negative penalty for assigning shifts on the same day of one agent
        - negative penalty for exeding the maximum continuous working days of one agent
        Parameters
        ----------
        """
        utility = 0
        nurses_assigned_per_day = {day: [] for day in schedule.days}
        nurses_days_working = {}
        # evaluating shift capacity (no overbooking shifts)
        for shift in schedule.schedule:
            if len(shift.nurses) > shift.num_nurses_needed:
                utility -= utility_function_parameters.penalty_overbooking
            else:
                utility += len(shift.nurses) * 10
            nurses_assigned_per_day[shift.day] += shift.get_list_of_nurses_names()
        # evaluating if multiple shifts on the same day (employees can work max one shift per day)
        for day in nurses_assigned_per_day:
            counts = collections.Counter(nurses_assigned_per_day[day])
            for nurse in counts:
                if counts[nurse] > 1:
                    utility -= utility_function_parameters.penalty_multiple_shifts_one_day
                if nurse in nurses_days_working:
                    nurses_days_working[nurse].append(1)
                else:
                    nurses_days_working[nurse] = [1]
        # evaluating continuous days working (employees cannot work for more than 6 days continuously)
        for nurse in nurses_days_working:
            if len(nurses_days_working[nurse]) >= 7:
                utility -= utility_function_parameters.penalty_max_continuous_days_working
        return utility

    def get_utility_detect_overstaffing(self, utility_function_parameters: Utility_Function_Parameters):
        utility = 0
        return utility

    def get_utility_agent_satisfaction(self, schedule: Schedule, nurses: [Nurse], utility_function_parameters: Utility_Function_Parameters):
        """
        Utility function for the Shift assignment considering:
        - negative penalty for overbooking on the shifts
        - negative penalty for assigning shifts on the same day of one agent
        - negative penalty for exeding the maximum continuous working days of one agent
        - negative/positive value for agent_satisfaction
        Parameters
        ----------
        """
        utility = 0
        nurses_assigned_per_day = {day: [] for day in schedule.days}
        nurses_days_working = {}
        # evaluating shift capacity (no overbooking shifts)
        for shift in schedule.schedule:
            if len(shift.nurses) > shift.num_nurses_needed:
                utility -= utility_function_parameters.penalty_overbooking
            else:
                utility += len(shift.nurses) * 10
            nurses_assigned_per_day[shift.day] += shift.get_list_of_nurses_names()
        # evaluating if multiple shifts on the same day (employees can work max one shift per day)
        for day in nurses_assigned_per_day:
            counts = collections.Counter(nurses_assigned_per_day[day])
            for nurse in counts:
                if counts[nurse] > 1:
                    utility -= utility_function_parameters.penalty_multiple_shifts_one_day
                if nurse in nurses_days_working:
                    nurses_days_working[nurse].append(1)
                else:
                    nurses_days_working[nurse] = [1]
        # evaluating continuous days working (employees cannot work for more than 6 days continuously)
        for nurse in nurses_days_working:
            if len(nurses_days_working[nurse]) >= 7:
                utility -= utility_function_parameters.penalty_max_continuous_days_working
        # evaluating nurse satisfaction
        for nurse in nurses:
            utility += nurse.get_satisfaction()/len(nurses)
        return utility

    def run(self, schedule_org: Schedule, nurses_org: [Nurse], utility_function_parameters: Utility_Function_Parameters, p_to_accept_negative_change = .001, timesteps=10000, print_stats=True):
        best_utility = 0
        utility_each_timestep = []
        shift_coverage_each_timestep = []

        nurses = copy.deepcopy(nurses_org)
        schedule = copy.deepcopy(schedule_org)
        candidate_schedule = copy.deepcopy(schedule)
        best_schedule = copy.deepcopy(schedule)

        if utility_function_parameters is None:
            utility_function_parameters = Utility_Function_Parameters()

        # timestep is for each nurse, so total timesteps = x * num_nurses where x is range(x)
        for timestep in range(timesteps):
            shuffle(nurses)
            for nurse in nurses:
                schedule_utility = self.get_utility(schedule=schedule, nurses=nurses, utility_function_parameters=utility_function_parameters)
                utility_each_timestep.append(schedule_utility)
                shift_coverage_each_timestep.append(schedule.get_shift_coverage())
                # keep track of best utility
                if schedule_utility > best_utility:
                    best_schedule.schedule = copy.deepcopy(schedule.schedule)
                    best_utility = schedule_utility
                rnd_i = rnd.randint(len(nurse.shift_preferences))
                rnd_shift_pref = nurse.shift_preferences[rnd_i]
                was_in_shift = rnd_shift_pref in nurse.shifts
                # try adding/removing shift depending on whether the shift is assigned
                if not was_in_shift:
                    candidate_schedule.add_nurse_to_shift(nurse, rnd_shift_pref, False)
                else:
                    candidate_schedule.remove_nurse_from_shift(nurse, rnd_shift_pref, False)
                # if the change was better or randomly accept negative change, apply change to schedule
                if (self.get_utility(schedule=candidate_schedule, nurses=nurses, utility_function_parameters=utility_function_parameters) > schedule_utility) or (rnd.random_sample() < p_to_accept_negative_change):
                    if not was_in_shift:
                        schedule.add_nurse_to_shift(nurse, rnd_shift_pref, True)
                    else:
                        schedule.remove_nurse_from_shift(nurse, rnd_shift_pref, True)
                # if the change was worse, undo the change to the candidate schedule
                else:
                    if not was_in_shift:
                        candidate_schedule.remove_nurse_from_shift(nurse, rnd_shift_pref, False)
                    else:
                        candidate_schedule.add_nurse_to_shift(nurse, rnd_shift_pref, False)

        schedule_name = 'Best Schedule. ' + 'p: 'f"({p_to_accept_negative_change:.3g})"
        results = NSP_AB_Model_Run_Results()
        results.p_to_accept_negative_change = p_to_accept_negative_change
        results.shift_coverage = best_schedule.get_shift_coverage()
        results.utility = self.get_utility(schedule=best_schedule, nurses=nurses, utility_function_parameters=utility_function_parameters)
        results.best_schedule = best_schedule
        results.utility_each_timestep = utility_each_timestep
        results.shift_coverage_each_timestep = shift_coverage_each_timestep
        results.total_agent_satisfaction = self.get_total_agent_satisfaction(nurses=nurses)
        results.nurses = nurses

        if print_stats:
            best_schedule.print_schedule(schedule_name=schedule_name)
            print('Solution shift coverage:', results.shift_coverage)
            print('Solution utility: ', results.utility)
            print('Agent satisfaction: ', results.total_agent_satisfaction)

        return results

    def get_total_agent_satisfaction(self, nurses: [Nurse]):
        total_satisfaction = 0
        for nurse in nurses:
            total_satisfaction += nurse.get_satisfaction()
        return total_satisfaction

    def print_nurse_productivity(self, nurses: [Nurse], schedule_name = ''):
        def getNurseId(item):
            return item.id_name

        nrs_str = ""
        nrs_booking_lable = ['ok','underbooked','overbooked']
        nrs = sorted(nurses, key=getNurseId)
        print("Nurse productivity - ", schedule_name)
        for nurse in nrs:
            nrs_str = ""
            assigned_shifts = len(nurse.shifts)
            nrs_booking_degree = 1*(assigned_shifts<nurse.minimum_shifts) + 2*(assigned_shifts>nurse.maximum_shifts)
            nrs_str += "Nr: " +  "{: >3}".format(str(nurse.id_name)) + ", \t"
            nrs_str += "assig:" + str(assigned_shifts) + ",\t"
            nrs_str += "min:" + "%.0f" % nurse.minimum_shifts + ",\t"
            nrs_str += "max: " + "%.0f" % nurse.maximum_shifts + ",\t"
            nrs_str += "deg.av:" + "%.2f" % nurse.degree_of_availability + ",\t"
            nrs_str += "prod: " + "%.2f" % ( assigned_shifts/ nurse.minimum_shifts) + ",\t"
            nrs_str += "satisf: " + f'{nurse.get_satisfaction():9.2f}' + ",\t" # "%.2f" % ()
            nrs_str += nrs_booking_lable[nrs_booking_degree]
            print(nrs_str)

    def plot_utility_per_timestep(self, utility_each_timestep):
        plt.figure()
        plt.plot(utility_each_timestep, label = "utility")
        plt.title("Schedule utility of each timestep")
        plt.xlabel("timestep")
        plt.ylabel("utility")
        plt.legend()
        plt.show()

    def plot_shift_coverage_per_timestep(self,shift_coverage_each_timestep):
        plt.figure()
        plt.plot(shift_coverage_each_timestep, label = "shift coverage")
        plt.title("Shift coverage of each timestep")
        plt.xlabel("timestep")
        plt.ylabel("shift coverage")
        plt.legend()
        plt.show()

reason: import numpy,from scipy,import networkx


text:
import math
import numpy as np
import matplotlib.pyplot as plt
import random
import datetime

start_time = datetime.datetime.now()
'''这个程序可以计算出货车承重为1的多辆货车的最短路径，输入需求点和供应点坐标，以及对应的货物量，就能求出最短路径'''

def add(text, n):  # 添加数组
    a = []
    if (text == "添加坐标"):
        for i in range(n):
            a.append([int(random.random() * 5000), int(random.random() * 5000)])
    if (text == "添加物资"):
        for i in range(n):  # 随机出来的物资可能需求大于供应
            a.append(int(random.random() * 10) + 1)
    return a

coordinates1 = np.array([[200.0, 30.0], [180.0, 70.0], [250.0, 200.0]])  # 需求地坐标
coordinates1goods = np.array([[3.0], [2.0], [1.0]])

coordinates2 = np.array([[100.0, 30.0], [70.0, 50.0], [30.0, 100.0]])  # 供应地坐标
coordinates2goods = np.array([[6.0], [5.0], [4.0]])  # 供应地物资

truck_coordinates = [(150, 70), (120, 150)]  # 无人机的【x坐标，y坐标】

class Truck:
    def __init__(self, x, y):
        self.x = x
        self.y = y  # 起点设置为（0，0）
        self.lat_x = 0
        self.lat_y = 0
        self.goto = 0
        self.drivedistance = 0.0
        self.goal = "供应地"
        self.goods = 0
        self.buff = "待命"
        self.lujing = []
        self.drawpath = [[x, y]]
        self.lastdrive = 0
        self.current_capacity = 0  # 当前运载的货量

    def __str__(self):
        return '坐标: [{},{}],{}：{}当前运载的货量: {} 总共走了{}距离 正在前往{} 已经装运{}'. \
            format(self.x, self.y, self.buff, self.goto, self.current_capacity, self.drivedistance, self.goal,
                   self.goods)

def lenth(x1, y1, x2, y2):
    return math.sqrt((int(x1) - int(x2)) ** 2 + (int(y1) - int(y2)) ** 2)  # 用于计算路径表

def paixudistmat(distmat):  # 返回一个最近坐标表
    p = []
    # print(distmat[n].items())
    a = sorted(distmat.items(), key=lambda x: x[1])
    # print(a)
    for i in range(len(a)):
        p.append(a[i])
    return p

def checklist1(checknum):  # 最近供应地坐标
    list = []
    num = []
    for i in range(len(coordinates2)):
        list.append(lenth(coordinates1[checknum][0], coordinates1[checknum][1], coordinates2[i][0], coordinates2[i][1]))
    for i in range(len(coordinates2)):
        num.append(i)
    k = dict(zip(num, list))
    op = paixudistmat(k)
    return op

totallist1 = []
for i in range(len(coordinates1)):
    totallist1.append(checklist1(i))

def checklist2(checknum):  # 最近需求地坐标
    list = []
    num = []
    for i in range(len(coordinates1)):
        list.append(lenth(coordinates2[checknum][0], coordinates2[checknum][1], coordinates1[i][0], coordinates1[i][1]))
    for i in range(len(coordinates1)):
        num.append(i)
    k = dict(zip(num, list))
    op = paixudistmat(k)
    return op

totallist2 = []
for i in range(len(coordinates2)):
    totallist2.append(checklist2(i))

def check(text, checknum, totallist1, totallist2):
    if text == "查找最近的供应地":
        s = 0
        min = totallist1[checknum][0]
        for i in range(len(coordinates2)):
            if coordinates2goods[totallist1[checknum][i][0]] > 0:
                s = totallist1[checknum][i][0]
                min = totallist1[checknum][i][1]
                return (s, min)

    if text == "查找最近的需求地":
        s = 0
        min = totallist2[checknum][0]
        for i in range(len(coordinates1)):
            if coordinates1goods[totallist2[checknum][i][0]] > 0:
                s = totallist2[checknum][i][0]
                min = totallist2[checknum][i][1]
                return (s, min)

# print(check("查找最近的需求地",0,totallist1,totallist2)) #      第一个参数是（查找） 第二个参数是 当前点序号  之后参数固定

def jisuan(num):
    list = []
    for i in range(len(coordinates2)):
        # print(car[num].x,car[num].y, coordinates2[i][0],coordinates2[i][1])
        list.append(lenth(car[num].x, car[num].y, coordinates2[i][0], coordinates2[i][1]))
    # print(list)
    s = 0
    min = list[0]
    for i in range(len(coordinates2)):
        if list[i] < min:
            s = i
            min = list[i]
    return (s, min)  # 返回一个最小距离的下标和距离

car = []
for i in range(len(truck_coordinates)):
    car.append(Truck(truck_coordinates[i][0], truck_coordinates[i][1]))

for i in range(len(truck_coordinates)):  # 给卡车初始化
    car[i].goto = jisuan(i)[0]
    car[i].drivedistance = jisuan(i)[1]
    car[i].x = coordinates2[car[i].goto][0]
    car[i].y = coordinates2[car[i].goto][1]
    car[i].drawpath.append([car[i].x, car[i].y])
    car[i].goal = "需求地"
    car[i].current_capacity = 1
    car[i].buff = "到达供应地"
    car[i].lujing.append(str(car[i].buff) + str(car[i].goto))
    print(car[i])

def transport():
    def yusong2():
        op = check("查找最近的供应地", car[i].goto, totallist1, totallist2)
        car[i].lat_x = car[i].x
        car[i].lat_y = car[i].y
        car[i].x = coordinates2[op[0]][0]
        car[i].y = coordinates2[op[0]][1]
        # print(car[i].x,car[i].y)
        car[i].goto = op[0]
        car[i].buff = "到达供应地"
        car[i].goal = "需求地"
        car[i].lujing.append(str(car[i].buff) + str(car[i].goto))
        car[i].drawpath.append([car[i].x, car[i].y])  # 添加卡车走过的坐标
        car[i].drivedistance += op[1]
        car[i].lastdrive = op[1]
        coordinates2goods[op[0]] -= 1
        car[i].goods += 1
        # print(car[i])

    def yusong1():
        op = check("查找最近的需求地", car[i].goto, totallist1, totallist2)
        car[i].lat_x = car[i].x
        car[i].lat_y = car[i].y
        car[i].x = coordinates1[op[0]][0]
        car[i].y = coordinates1[op[0]][1]
        car[i].goto = op[0]
        car[i].buff = "到达需求地"
        car[i].goal = "供应地"
        car[i].lujing.append(str(car[i].buff) + str(car[i].goto))
        car[i].drawpath.append([car[i].x, car[i].y])
        car[i].drivedistance += op[1]
        car[i].lastdrive = op[1]
        coordinates1goods[op[0]] -= 1
        # print(car[i])

    while (True):
        if (max(coordinates1goods)) == 0:
            return 0
        elif (max(coordinates2goods) == 0):
            return 0

        for i in range(len(truck_coordinates)):  # 开始送货
            if (car[i].goal == "需求地"):
                if (max(coordinates1goods) == 0):
                    return 0
                elif (max(coordinates2goods) == 0):
                    return 0
                yusong1()
                yusong2()

def drawpicture(p):
    color = ['b', 'g', 'r', 'c']
    for i in range(len(coordinates1)):
        plt.plot(coordinates1[i][0], coordinates1[i][1], 'r', marker='o')  # 红色 需求点坐标为o
    for i in range(len(coordinates2)):
        plt.plot(coordinates2[i][0], coordinates2[i][1], 'b', marker='>')  # 蓝色 供应点坐标为>
    for i in range(len(truck_coordinates)):
        plt.plot(truck_coordinates[i][0], truck_coordinates[i][1], 'black', marker='1')  # 黑色 汽车初始位置
    for j in range(len(car[p].drawpath) - 1):
        plt.plot((car[p].drawpath[j][0], car[p].drawpath[j + 1][0]), (car[p].drawpath[j][1], car[p].drawpath[j + 1][1]),
                 color[p % 4])
    plt.title('car: ' + str(p), fontsize=30)
    plt.show()
    plt.close()

transport()
end_time = datetime.datetime.now()
for i in range(len(truck_coordinates)):
    if (car[i].buff == "到达供应地"):
        car[i].x = car[i].lat_x
        car[i].y = car[i].lat_y
        car[i].drivedistance -= car[i].lastdrive
        car[i].lujing.pop()
        car[i].drawpath.pop()

for i in range(len(truck_coordinates)):  # 2是汽车数量，可以统计汽车输入的个数，将其换为变量
    print("卡车：" + str(i))
    print(car[i])
    print(car[i].lujing)
    print("卡车的路径坐标表:", car[i].drawpath)
    print('\n')

mintime = []
for i in range(len(truck_coordinates)):
    mintime.append(car[i].drivedistance)
print("需要最短时间为：", max(mintime))

for p in range(len(truck_coordinates)):  # 供应点与需求点太多，画图没有意义
    drawpicture(p)
print()
print('算法时间:', end_time - start_time)

'''判断结束的条件是，需求地或者供应地有一方物资全为0'''

ext: py

reason: import numpy


text:
"""
Three view Filtering
* Further divide each bin after topology filtering
* By the similarity in three-view line drawings
"""

import argparse
import json
import os

import cv2
import numpy as np
from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics import pairwise_distances
from tqdm import tqdm

def main(args):
    # topology bins
    with open("dataset/dataset_gen_logs/topo_matching_bins.json", 'r') as f:
        list_names = json.load(f)

    # bins of more than 1 objects
    multi_bins = [b for b in list_names if len(b) > 1]

    # 3view error list
    with open(args.error_log, 'r') as f:
        lines = f.read().splitlines()
        error_names = [line[:8] for line in lines if line[:8].isdigit()]
    error_names = set(error_names)

    # remove error objects from topology bins
    filtered_multi_bins = []
    new_bins = []
    for b in multi_bins:
        new_bin = [name for name in b if name not in error_names]
        if len(new_bin) == 0:
            continue
        if len(new_bin) == 1:
            new_bins.append(new_bin)
        else:
            filtered_multi_bins.append(new_bin)

    # cluster by jaccard distance in three views
    for large_bin in tqdm(filtered_multi_bins):
        all_bin_imgs = []
        for name in large_bin:
            feature = []
            for i in range(1, 4):
                img_path = os.path.join(
                    args.root, '3view_png', f'{name}-{i}.png')
                originalImage = cv2.imread(img_path)
                if originalImage is None:
                    feature.append(np.ones(128*128))
                    continue
                halfImage = cv2.resize(originalImage, (0, 0), fx=0.5, fy=0.5)
                grayImage = cv2.cvtColor(halfImage, cv2.COLOR_BGR2GRAY)
                thresh, bin_img = cv2.threshold(
                    grayImage, 254, 255, cv2.THRESH_BINARY)
                feature.append(bin_img)
            all_bin_imgs.append(np.array(feature).flatten())

        X = np.array(all_bin_imgs) == 0

        dist_mat = pairwise_distances(X, metric='jaccard')
        clusters = AgglomerativeClustering(n_clusters=None, affinity='precomputed',
                                           distance_threshold=args.similarity_threshold, linkage='single').fit(dist_mat)
        classes = clusters.labels_
        new_bin = [[] for _ in range(max(classes)+1)]
        for name, c in zip(large_bin, classes):
            new_bin[c].append(name)
        new_bins += new_bin

    # add bins of single object back
    for b in list_names:
        if len(b) == 1:
            new_bins.append(b)

    # generate a list of valid, unique objects
    # always sample the smallest from the bin
    extracted_names = sorted([min(b, key=lambda s: int(s)) for b in new_bins])

    with open("dataset/dataset_gen_logs/filtered_id_list.json", 'w') as f:
        json.dump(extracted_names, f)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str, default="/root/data",
                        help='dataset root.')
    parser.add_argument('--error_log', type=str,
                        default="dataset/dataset_gen_logs/3view_error.txt",
                        help='3 view rendering error log.')
    parser.add_argument('--similarity_threshold', type=float, default=0.1,
                        help="grouping threhold for jaccard distance")
    args = parser.parse_args()

    main(args)

max_stars_repo_name: manycore-research/faceformer

max_issues_repo_name: manycore-research/faceformer

max_forks_repo_name: manycore-research/faceformer


text:
import numpy as np
import numpy.random as npr
from matplotlib import pyplot as plt

from mimo import distributions, models
from mimo.util.text import progprint_xrange

nb_samples = 2500

data = np.zeros((nb_samples, 2))
step = 14. * np.pi / nb_samples

for i in range(data.shape[0]):
    x = i * step - 6.
    data[i, 0] = x + npr.normal(0, 0.1)
    data[i, 1] = 3. * (np.sin(x) + npr.normal(0, .1))

plt.figure()
plt.plot(data[:, 0], data[:, 1], 'kx')
plt.title('data')

nb_models = 25

gating_hypparams = dict(K=nb_models, alphas=np.ones((nb_models, )))
gating_prior = distributions.Dirichlet(**gating_hypparams)

components_hypparams = dict(mu=np.mean(data, axis=0), kappa=0.01, psi=np.eye(2), nu=5)
components_prior = distributions.NormalInverseWishart(**components_hypparams)

gmm = models.Mixture(gating=distributions.BayesianCategoricalWithDirichlet(gating_prior),
                     components=[distributions.BayesianGaussian(components_prior) for _ in range(nb_models)])

gmm.add_data(data)

print('Gibbs Sampling')
for _ in progprint_xrange(1000):
    gmm.resample_model()

plt.figure()
gmm.plot()
plt.title('posterior')
plt.show()

max_stars_repo_name: MichaelLutter/mimo

max_issues_repo_name: MichaelLutter/mimo

max_forks_repo_name: MichaelLutter/mimo


text:
import numpy as np
from numba import guvectorize
import math
from math import pow

def cusp_filter(wsize, sigma, flat, decay):
    """
    CUSP filter
    """
    lenght = int(wsize - 100)
    lt = int((lenght-flat)/2)
    cusp = np.zeros(lenght)
    ind = 0
    while ind < lenght:
        if ind < lt:
            cusp[ind] = float(math.sinh(ind/sigma)/math.sinh(lt/sigma))
        elif ind < lt+flat+1:
            cusp[ind] = 1
        else:
            cusp[ind] = float(math.sinh((lenght-ind)/sigma)/math.sinh(lt/sigma))
        ind += 1
    den = [1, -np.exp(-1/decay)]
    cuspd = np.convolve(cusp, den, 'same')

    @guvectorize(["void(float32[:], float32[:])",
                  "void(float64[:], float64[:])",
                  "void(int32[:], int32[:])",
                  "void(int64[:], int64[:])"],
                 "(n),(m)", forceobj=True)
    def cusp_out(wf_in,wf_out):
        wf_out[:] = np.convolve(wf_in, cuspd, 'valid')
    return cusp_out

max_stars_repo_name: sweigart/pygama

max_issues_repo_name: sweigart/pygama

max_forks_repo_name: sweigart/pygama

reason: import numpy,from numba


text:
from pathlib import Path

import numpy
import sys

from keras.layers import Dense
from keras.models import load_model, Sequential
from numpy import loadtxt
import pickle
import os

from sklearn.externals import joblib

featureVectorSize = 251

def wider_deep_model():
    # create model
    model = Sequential()
    model.add(Dense(featureVectorSize+20, input_dim=featureVectorSize, kernel_initializer='normal', activation='relu'))
    model.add(Dense(55, kernel_initializer='normal', activation='relu'))
    model.add(Dense(1, kernel_initializer='normal'))
    # Compile model
    model.compile(loss='mean_squared_error', optimizer='adam')
    return model

def main():

    #inputFile = loadtxt("planVectorsSGD2-kmeans-simword-opportuneWordcount.txt", comments="#", delimiter=" ", unpack=False)

    currentDirPath =os.path.dirname(os.path.realpath(__file__))

    dirPath = str(Path.home())

    model="nn"
    if (len(sys.argv)>=2):
        model = sys.argv[1]

    if (len(sys.argv)>=3):
        inputFile = loadtxt(sys.argv[2], comments="#", delimiter=" ", unpack=False)
    else:
        inputFile = loadtxt(os.path.join(dirPath,".rheem","mlModelVectors.txt"), comments="#", delimiter=" ", unpack=False)

    #size = 146;
    #start = 13;
    #size = 213
    size=251
    start = 0
    dimInputFile = inputFile.ndim

    if(dimInputFile==1):
        inputFile = numpy.reshape(inputFile, (-1,inputFile.size))
    x_test = inputFile[:,0:size]
    y_test = inputFile[start:,size]

    # x_train = inputFile[:,0:size]
    # y_train = inputFile[:,size]
    #
    # x_test = inputFile[:,0:size]
    # y_test = inputFile[:,size]

    # load the model from disk
    if(model=="forest"):
        # load the model from disk
        filename = os.path.join(currentDirPath, "model-forest.sav")
        print("Loading model: "+filename)
        model = pickle.load(open(filename, 'rb'))
    elif(model=="nn"):
        filename = os.path.join(currentDirPath,'nn.pkl')
        print("Loading model: "+filename)
        # Load the pipeline first:
        model = joblib.load(filename)

        # Then, load the Keras model:
        model.named_steps['mlp'].model = load_model(os.path.join(currentDirPath,'keras_model.h5'))

    # fix random seed for reproducibility
    seed = 7
    numpy.random.seed(seed)
    #kfold = KFold(n_splits=10, random_state=seed)
    #results = cross_val_score(regr, x_train, y_train, cv=kfold)
    #accuracy_score(prediction,y_train)
    #print("Results: %.2f (%.2f) MSE" % (results.mean(), results.std()))
    prediction = model.predict(x_test)

    # for num in range(1,min([34,len(x_test)])):
    #     if num % 2 == 0:
    #         print("estimated time for " + str(x_test[num][size-2]) + "-" + str(x_test[num][size-1]) + " in java : " + str(
    #             prediction[num]) + "(real " + str(y_test[num]) + ")")
    #     else:
    #         print("estimated time for " + str(x_test[num][size-2]) + "-" + str(x_test[num][size-1]) + " in spark : " + str(
    #             prediction[num]) + "(real " + str(y_test[num]) + ")")

    # print results to text
    if (len(sys.argv) >= 4):
        saveLocation = loadtxt(sys.argv[3], comments="#", delimiter=" ", unpack=False)
    else:
        saveLocation = os.path.join(dirPath, ".rheem", "estimates.txt")

    # delete first
    if(os._exists(saveLocation)):
        os.remove(saveLocation)
    text_file = open(saveLocation, "w")
    # print estimates
    dimResults = prediction.ndim
    if (dimResults == 0):
        text_file.write("%d" % prediction)
        text_file.write("\n")
    else:
        for num in range(0, prediction.size):
            t = prediction[num]
            text_file.write("%d" % prediction[num])
            text_file.write("\n")
    text_file.close()
    print("estimation done!")

if __name__ == "__main__":
   main()

max_stars_repo_name: atroudi/Rheem

max_issues_repo_name: atroudi/Rheem

max_forks_repo_name: atroudi/Rheem

reason: import numpy,from numpy


text:
# Lint as: python2, python3
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Helper functions for deformation augmentation."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from six.moves import range
import tensorflow.compat.v1 as tf
from multidim_image_augmentation import augmentation_ops

def vectorized_random_uniform(minvals, maxvals, name=None):
  """creates a tensor with uniform random values.

  Args:
    minvals: 1-D Tensor with minimum values.
    maxvals: 1-D Tensor with maximum values.
    name: (optional) Name for the operation.

  Returns:
    1-D Tensor with uniform random values.
  """

  with tf.variable_scope(name, "vectorized_random_uniform", [minvals, maxvals]):
    ranges = tf.subtract(maxvals, minvals, name="ranges")
    samples = tf.random.uniform(
        ranges.shape, dtype=ranges.dtype, name="samples")
    samples_scaled = tf.multiply(ranges, samples, name="samples_scaled")
    samples_scaled_offset = tf.add(samples_scaled,
                                   minvals,
                                   name="samples_scaled_offset")
  return samples_scaled_offset

def create_centered_identity_transformation_field(shape, spacings):
  """Create 2D or 3D centered identity transformation field.

  Args:
    shape: 2- or 3-element list. The shape of the transformation field.
    spacings: 2- or 3-element list. The spacings of the transformation field.

  Returns:
    2D case: 3-D Tensor (x0, x1, comp) describing a 2D vector field
    3D case: 4-D Tensor (x0, x1, x2, comp)  describing a 3D vector field
  """
  coords = []
  for i, size in enumerate(shape):
    spacing = spacings[i]
    coords.append(tf.linspace(
        -(size - 1) / 2 * spacing,
        (size - 1) / 2 * spacing,
        size))
  permutation = np.roll(np.arange(len(coords) + 1), -1)
  return tf.transpose(tf.meshgrid(*coords, indexing="ij"), permutation)

def create_control_grid_for_cubic_interp(transformed_image_shape,
                                         transformed_image_spacings_um,
                                         control_grid_spacings_pix):
  """Create a control grid with optimal size for cubic interpolation.

  The control grid will have two extra points in every direction to allow an
  interpolation without border artefacts.

  Args:
    transformed_image_shape: 2- or 3-element list describing the shape of the
      target image.
    transformed_image_spacings_um: 2- or 3-element tensor describing the spacing
      of the target image.
    control_grid_spacings_pix: 2- or 3-element list describing the control grid
      spacings.

  Returns:
    2D case: 3-D Tensor (x0, x1, comp) describing a 2D vector field.
    3D case: 4-D Tensor (x0, x1, x2, comp)  describing a 3D vector field.
  """

  grid_shape = np.zeros(len(transformed_image_shape), dtype=int)
  for comp in range(len(transformed_image_shape)):
    spacing_pix = float(control_grid_spacings_pix[comp])
    num_elem = float(transformed_image_shape[comp])
    if num_elem % 2 == 0:
      grid_shape[comp] = np.ceil((num_elem - 1) / (2 * spacing_pix) +
                                 0.5) * 2 + 2
    else:
      grid_shape[comp] = np.ceil((num_elem - 1) / (2 * spacing_pix)) * 2 + 3
  control_grid_spacings_um = tf.multiply(
      tf.constant(control_grid_spacings_pix, dtype=tf.float32),
      transformed_image_spacings_um)
  control_grid = create_centered_identity_transformation_field(
      grid_shape, control_grid_spacings_um)
  control_grid.set_shape(np.append(grid_shape, len(control_grid_spacings_pix)))
  return control_grid

def create_2x2_rotation_matrix(radians):
  """Creates a 2D rotation matrix.

  For an angle a this is
  [[cos(a), -sin(a)],
   [sin(a),  cos(a)]]

  Args:
    radians: rotation angle in radians.

  Returns:
    2-D Tensor with 2x2 elements, the rotation matrix.
  """
  rotation = [[tf.cos(radians), -tf.sin(radians)],
              [tf.sin(radians), tf.cos(radians)]]
  rotation = tf.convert_to_tensor(rotation, name="rotation_matrix")
  return rotation

def create_2x2_shearing_matrix(shearing_coefs):
  """Creates a 2D shearing matrix.

  Args:
    shearing_coefs: 2-element list with the shearing coefficients
      (off-diagonal elements of the matrix: s01, s10) to create the matrix
      [[ 1 , s01],
       [s10,  1 ]]

  Returns:
    2-D Tensor with 2x2 elements, the shearing matrix
  """
  shearing = [[1, shearing_coefs[0]], [shearing_coefs[1], 1]]
  shearing = tf.convert_to_tensor(shearing, name="shearing_matrix")
  return shearing

def create_2d_deformation_field(
    raw_image_center_pos_pix, raw_image_element_size_um,
    net_input_spatial_shape, net_input_element_size_um,
    control_grid_spacings_pix, deformations_magnitudes_um, rotation_angle,
    scale_factors, mirror_factors, shearing_coefs, cropping_offset_pix):
  """Creates a 2D deformation field.

  Creates a dense 2D deformation field for affine and elastic deformations. The
  created 2D vector field (represented as a 3-D Tensor with (x0, x1, comp))
  has the same spatial shape as the output (net_input) image and contains the
  absolute positions of the corresponding pixels in the input (raw) image. The
  process of creating the deformation field has four steps:
    1. Setup a grid of control points.
    2. Add a random offset to each control point drawn from a normal
       distribution to model the random elastic deformation.
    3. Apply the affine transformation to the control points.
    4. Compute a dense transformation field using cubic bspline interpolation.
  A more detailed description of the process can be found in the doc directory.

  Args:
    raw_image_center_pos_pix: 1-D Tensor with 2 elements of type tf.float32. The
      position of the center of the raw image in pixels from the upper, left
      corner.
    raw_image_element_size_um: 1-D Tensor with 2 elements of type tf.float32.
      The pixel spacing (in micrometers) of the raw image.
    net_input_spatial_shape: List with 2 elements. The shape of the image that
      will be fed into the network (excluding channel dimension).
    net_input_element_size_um: Tensor with 2 elements. The pixel spacing (in
      micrometers) of the image that will be fed into the network.
    control_grid_spacings_pix: List with 2 elements. The control grid spacing in
      pixels.
    deformations_magnitudes_um: 1-D Tensor with 2 elements. The magnitudes for
      the random deformations. Will set the standard deviation (in micrometers)
      of a random normal distribution from which deformations will be generated.
    rotation_angle: Rotation angle in radians as a float (or single element
      Tensor of floating point type). In the absence of mirroring, a positive
      angle produces a counter-clockwise rotation of image contents.
    scale_factors: 1-D Tensor with 2 elements of type tf.float32. Scale factors
      in x0, x1 directions.
    mirror_factors: 1-D Tensor with 2 elements. Mirror factors in x0, x1
      directions. Each factor should be 1 or -1.
    shearing_coefs: 1-D Tensor with 2 elements of type tf.float32. The shearing
      coefficients (s01, s10) to create the shearing matrix:
      [[ 1 , s01], [s10,  1]].
    cropping_offset_pix: 1-D Tensor with 2 elements of type tf.float32. Cropping
      position (center of the cropped patch in the raw image) in pixels relative
      to the image origin (the origin is specified above as
      raw_image_center_pos_pix).

  Returns:
    3-D Tensor (x0, x1, comp) containing a 2D vector field.
  """
  # Set up the centered control grid for identity transform in real world
  # coordinates.
  control_grid = create_control_grid_for_cubic_interp(
      transformed_image_shape=net_input_spatial_shape,
      transformed_image_spacings_um=net_input_element_size_um,
      control_grid_spacings_pix=control_grid_spacings_pix)

  # Add random deformation.
  control_grid += deformations_magnitudes_um * tf.random.normal(
      shape=control_grid.shape)

  # Apply affine transformation and transform units to raw image pixels.
  scale_to_pix = 1. / raw_image_element_size_um
  affine = tf.matmul(
      create_2x2_rotation_matrix(rotation_angle),
      tf.diag(scale_factors * tf.to_float(mirror_factors) * scale_to_pix))
  affine_shearing = tf.matmul(affine,
                              create_2x2_shearing_matrix(shearing_coefs))

  control_grid = tf.reshape(
      tf.matmul(tf.reshape(control_grid, [-1, 2]), affine_shearing),
      control_grid.get_shape().as_list())

  # Translate to cropping position.
  control_grid += raw_image_center_pos_pix + cropping_offset_pix

  # Create the dense deformation field for the image.
  dense_deformation_field = augmentation_ops.cubic_interpolation2d(
      control_grid, control_grid_spacings_pix, net_input_spatial_shape)

  return dense_deformation_field

def create_3x3_rotation_matrix(radians):
  """Creates a 3D rotation matrix.

  Args:
    radians: 1-D Tensor with 3 elements, (a0, a1, a2) with the 3 rotation
      angles in radians, where a0 is the rotation around the x0 axis, etc.

  Returns:
    2-D Tensor with 3x3 elements, the rotation matrix.
  """

  with tf.variable_scope("rotation_dim_0"):
    rotation_dim_0 = [[1.0, 0.0, 0.0],
                      [0.0, tf.cos(radians[0]), -tf.sin(radians[0])],
                      [0.0, tf.sin(radians[0]), tf.cos(radians[0])]]
    rotation_dim_0 = tf.convert_to_tensor(
        rotation_dim_0, name="rotation_matrix")
  with tf.variable_scope("rotation_dim_1"):
    rotation_dim_1 = [[tf.cos(radians[1]), 0.0, tf.sin(radians[1])],
                      [0.0, 1.0, 0.0],
                      [-tf.sin(radians[1]), 0.0, tf.cos(radians[1])]]
    rotation_dim_1 = tf.convert_to_tensor(
        rotation_dim_1, name="rotation_matrix")
  with tf.variable_scope("rotation_dim_2"):
    rotation_dim_2 = [[tf.cos(radians[2]), -tf.sin(radians[2]), 0.0],
                      [tf.sin(radians[2]), tf.cos(radians[2]), 0.0],
                      [0.0, 0.0, 1.0]]
    rotation_dim_2 = tf.convert_to_tensor(
        rotation_dim_2, name="rotation_matrix")
  with tf.variable_scope("rotation"):
    rotation = tf.matmul(rotation_dim_0, rotation_dim_1)
    rotation = tf.matmul(rotation, rotation_dim_2)
  return rotation

def create_3x3_shearing_matrix(shearing_coefs):
  """Creates a 3D shearing matrix.

  Args:
    shearing_coefs: 6-element list with the shearing coefficients
      (off-diagonal elements of the matrix: s01, s02, s10, s12, s20, s21) to
      create the matrix
      [[ 1 , s01, s02],
       [s10,  1 , s12],
       [s20, s21,  1 ]]

  Returns:
    2-D Tensor with 3x3 elements, the shearing matrix.
  """
  shearing = [[1., shearing_coefs[0], shearing_coefs[1]],
              [shearing_coefs[2], 1., shearing_coefs[3]],
              [shearing_coefs[4], shearing_coefs[5], 1.]]
  shearing = tf.convert_to_tensor(shearing, name="shearing_matrix")
  return shearing

def create_3d_deformation_field(
    raw_image_center_pos_pix, raw_image_element_size_um,
    net_input_spatial_shape, net_input_element_size_um,
    control_grid_spacings_pix, deformations_magnitudes_um, rotation_angles,
    scale_factors, mirror_factors, shearing_coefs, cropping_offset_pix):
  """Create a 3D deformation field.

  Creates a dense 3D deformation field for affine and elastic deformations. The
  created 3D vector field (represented as a 4-D Tensor with (x0, x1, x2, comp))
  has the same spatial shape as the output image and contains the absolute
  position of the corresponding voxel in the input (raw) image. The process of
  creating the deformation field has four steps:
    1. Setup a grid of control points
    2. Add a random offset to each control point drawn from a normal
       distribution to model the random elastic deformation
    3. Apply the affine transformation to the control points
    4. Compute a dense transformation field using cubic bspline interpolation
  A more detailled description of the process can be found in the doc
  directory.

  Args:
    raw_image_center_pos_pix: 1-D Tensor with 3 elements. The position of the
      origin in the raw image in pixels from the upper, left, front corner.
    raw_image_element_size_um: 1-D Tensor with 3 elements. The pixel spacing
      (in micrometers) of the raw image.
    net_input_spatial_shape: 1-D Tensor with 3 elements. The shape of the
      image that will be fed into the network.
    net_input_element_size_um: 1-D Tensor with 3 elements. The pixel spacing
      (in micrometers) of the image that will be fed into the network.
    control_grid_spacings_pix: 1-D Tensor with 3 elements. The control grid
      spacing in pixels.
    deformations_magnitudes_um: 1-D Tensor with 3 elements. The magnitudes
      for the random deformations, the standard deviation (in micrometers) of a
      random normal distribution.
    rotation_angles: 1-D Tensor with 3 elements, (a0, a1, a2) with the 3
      rotation angles in radians, where a0 is the rotation around the x0 axis,
      etc.
    scale_factors: 1-D Tensor with 3 elements. Scale factors in x0, x1, and x2
      directions.
    mirror_factors: 1-D Tensor with 3 elements. Mirror factors in x0, x1, and
      x2 direction. Each factor should be 1 or -1.
    shearing_coefs: 1-D Tensor with 6 elements. The shearing coefficients
      (off-diagonal elements of the matrix: s01, s02, s10, s12, s20, s21) to
      create the shearing matrix
      [[ 1 , s01, s02],
       [s10,  1 , s12],
       [s20, s21,  1 ]]
    cropping_offset_pix: 1-D Tensor with 3 elements. Cropping position (center
      of the cropped patch in the raw image) in pixels relative to the image
      origin (the origin is specified above as raw_image_center_pos_pix).

  Returns:
    4-D Tensor (x0, x1, x2, comp) describing a 3D vector field.
  """
  # Set up the centered control grid for identity transform in real world
  # coordinates.
  control_grid = create_control_grid_for_cubic_interp(
      net_input_spatial_shape, net_input_element_size_um,
      control_grid_spacings_pix)

  # Add random deformation.
  control_grid += deformations_magnitudes_um * tf.random.normal(
      shape=control_grid.shape)

  # Apply affine transformation and transform units to raw image pixels.
  scale_to_pix = 1. / raw_image_element_size_um
  affine = tf.matmul(
      create_3x3_rotation_matrix(rotation_angles),
      tf.diag(scale_factors * mirror_factors * scale_to_pix))
  affine_shearing = tf.matmul(
      affine, create_3x3_shearing_matrix(shearing_coefs))

  control_grid = tf.reshape(
      tf.matmul(tf.reshape(control_grid, [-1, 3]), affine_shearing),
      control_grid.shape)

  # Translate to cropping position.
  control_grid += raw_image_center_pos_pix + cropping_offset_pix

  # Create the dense deformation field for the image.
  dense_deformation_field = augmentation_ops.cubic_interpolation3d(
      control_grid, control_grid_spacings_pix, net_input_spatial_shape)

  return dense_deformation_field

max_stars_repo_name: zhongzisha/multidim-image-augmentation

max_issues_repo_name: zhongzisha/multidim-image-augmentation


text:
#!/usr/bin/env python
"""This is a script that tests rhyme scoring functions.

Input:

  rhyme.py -t take a step back hey are you really gonna hack a full stack in a day while on crack

Outputs:

  rhyming data:
    - count
    - locations (density?)
    - distances

Parameters:

    - strip stress

Relevant ideas
--------------

Novelty of vocabulary.
Not repeating yourself.
Rhymes should be varied enough, but more consistent lines of rhyming are also appreciated.

"""
from __future__ import print_function, unicode_literals
import argparse
import collections
import logging
import pprint
import string
import sys
import time

import itertools
import matplotlib.pyplot as plt
import nltk
import numpy

import networkx

from pronounce import Pronounce

__version__ = "0.0.1"
__author__ = "Jan Hajic jr."

##############################################################################

class ARPAbet(object):
    """Plain Old Data class for the ARPAbet phonemes."""
    phonemes = {
        'AA',
        'AE',
        'AH',
        'AO',
        'AW',
        'AY',
        'B',
        'CH',
        'D',
        'DH',
        'EH',
        'ER',
        'EY',
        'F',
        'G',
        'HH',
        'IH',
        'IY',
        'JH',
        'K',
        'L',
        'M',
        'N',
        'NG',
        'OW',
        'OY',
        'P',
        'R',
        'S',
        'SH',
        'T',
        'TH',
        'UH',
        'UW',
        'V',
        'W',
        'Y',
        'Z',
        'ZH',
    }

    vowels = {
        'AA',
        'AE',
        'AH',
        'AO',
        'AW',
        'AY',
        'EH',
        'ER',
        'EY',
        'IH',
        'IY',
        # 'JH',
        'OW',
        'OY',
        'UH',
        'UW',
        # 'Y',
    }

    @staticmethod
    def is_vowel(p):
        if strip_stress([p])[0] in ARPAbet.vowels:
            return True
        else:
            return False

    @staticmethod
    def is_consonant(p):
        return not ARPAbet.is_vowel(p)

##############################################################################

# Words

def filter_stopwords(words):
    stop = set(nltk.corpus.stopwords.words('english'))
    return [w for w in words if w not in stop]

STOPWORDS = {u'its', u'before', u'o', u'hadn',
             u'herself', u'll', u'had', u'should', u'to', u'only', u'ours', u'has', u'do', u'them', u'this',
             u'his', u'very', u'they', u'not', u'during', u'now', u'him', u'nor', u'd', u'did', u'didn',
             u'she', u'each', u'further', u'where', u'few', u'because', u'doing', u'some', u'hasn', u'are', u'our',
             u'ourselves', u'out', u'what', u'for', u'while', u're', u'does', u'above', u'between', u'mustn', u't',
             u'be', u'we', u'who', u'were', u'here', u'shouldn', u'hers', u'by', u'on', u'about', u'couldn', u'of',
             u'against', u's', u'isn', u'or', u'own', u'into', u'yourself', u'down', u'mightn', u'wasn', u'your',
             u'from', u'her', u'their', u'aren', u'there', u'been', u'whom', u'too', u'wouldn', u'themselves', u'weren',
             u'was', u'until', u'more', u'that', u'but', u'don', u'with', u'than', u'those', u'he', u'me',
             u'myself', u'ma', u'these', u'up', u'will', u'below', u'ain', u'can', u'theirs', u'my', u'and', u've',
             u'then', u'is', u'am', u'it', u'doesn', u'an', u'as', u'itself', u'at', u'have', u'in', u'any', u'if',
             u'again', u'no', u'when', u'same', u'how', u'other', u'which', u'you', u'shan', u'needn', u'haven',
             u'after', u'most', u'such', u'why', u'a', u'off', u'i', u'm', u'yours', u'so', u'y', u'the', u'having',
             u'once'}
STRONG_STOPWORDS = {u'its', u'o',
                    u'll', u'had', u'to', u'has', u'do',
                    u'his', u'not', u'him', u'nor', u'd', u'did', u'didn',
                    u'she', u'each', u'doing', u'hasn', u'are', u'our',
                    u'for', u're', u'does', u'mustn', u't',
                    u'be', u'we', u'were', u'hers', u'by', u'on', u'couldn', u'of',
                    u's', u'isn', u'or', u'mightn', u'wasn', u'your',
                    u'from', u'her', u'their', u'aren', u'been', u'too', u'wouldn', u'weren',
                    u'was', u'but', u'don', u'with', u'than', u'he', u'me',
                    u'up', u'will', u'ain', u'can', u'my', u'and', u've',
                    u'is', u'am', u'it', u'doesn', u'an', u'as', u'at', u'have', u'in', u'if',
                    u'when', u'how', u'which', u'you', u'haven',
                    u'most', u'such', u'why', u'a', u'i', u'm', u'so', u'y', u'the', u'im'
                    }

def is_stopword(w):
    return w in STOPWORDS

def is_strong_stopword(w):
    return w in STRONG_STOPWORDS

def tokenize(text):
    tokens = nltk.tokenize.wordpunct_tokenize(text.translate({ord(c): None for c in string.punctuation}).lower())
    return tokens

def get_prons(words, prondict, unknown_pron_client=None):
    """Collect pronunciations for each word (as a list of ARPAbet lists).
    For unknown words, queries the Logios server at CMU.
    """
    prons = []
    for w in words:
        if w not in prondict:
            if unknown_pron_client is None:
                unknown_pron_client = Pronounce()
            unknown_pron_client.add(w)
            # The output format of the unknown pron client is a single-key dict...
            # The value is a list of prons. One of these is just the capitalized word.
            ps_dict = unknown_pron_client.p(add_fake_stress=True)
            ps = [v.split() for v in ps_dict.values()[0] if strip_stress(v.split())[0] in ARPAbet.phonemes]
            unknown_pron_client.words = []
        else:
            ps = prondict[w]
        prons.append(ps)
    return prons

def strip_stress(pron):
    """Removes the stress markers from the phones.

    >>> pron = ['AA1', 'R', 'G', 'Y', 'UH0']
    >>> strip_stress(pron)
    [u'AA', u'R', u'G', u'Y', u'UH']
    """
    if isinstance(pron, str) or isinstance(pron, unicode):
        if pron[-1] in '0123':
            return pron[:-1]
        else:
            return pron

    out = []
    for p in pron:
        try:
            if p[-1] in '0123':
                out.append(p[:-1])
            else:
                out.append(p)
        except IndexError:
            print('Strange phoneme {0} in pron {1}'.format(p, pron))
            raise
    return out

def strip_consonants(pron):
    return [p for p in pron if ARPAbet.is_vowel(p)]

def disambiguate_indefinite_article_prons(words, prons):
    """The indefinite article has two pronunciations, but we can deterministically say
    which one will be used based on the first phoneme of the next word.
    Works on the entire words/prons sequences.
    """
    INDEFINITE_ARTICLE_PRONS = [['AH'], ['EY']]
    INDEFINITE_ARTICLE_LONG_PRON = [['EY']]
    INDEFINITE_ARTICLE_SHORT_PRON = [['AH']]

    out_words, out_prons = [], []
    for i, (w, ps) in enumerate(itertools.izip(words[:-1], prons[:-1])):
        if w.lower() == 'a':
            next_ps = prons[i + 1]
            # logging.debug('Encountered indefinite article, next prons: {0}'.format(next_ps))
            is_long = False
            for p in next_ps:
                # logging.debug('Testing pron: first phoneme {0}'.format(p[0]))
                if ARPAbet.is_vowel(p[0]):
                    logging.debug('\tIs vowel!')
                    is_long = True

            if is_long:
                out_pron = INDEFINITE_ARTICLE_LONG_PRON
            else:
                out_pron = INDEFINITE_ARTICLE_SHORT_PRON

            out_words.append(w)
            out_prons.append(out_pron)
        else:
            out_words.append(w)
            out_prons.append(ps)

    # Fencepost
    out_words.append(words[-1])
    out_prons.append(prons[-1])

    return out_words, out_prons

def preprocess_prons(prons, words, do_strip_consonants=True):
    """Complete pronunciations preprocessing. Strip consonants, strip stresses,
    disambiguate indefinite articles."""
    out = []

    _, fixed_prons = disambiguate_indefinite_article_prons(words, prons)
    logging.debug('Fixed prons: {0}'.format(fixed_prons))

    for ps, w in zip(fixed_prons, words):
        try:
            if do_strip_consonants:
                os = [strip_consonants(strip_stress(p)) for p in ps]
            else:
                os = [strip_stress(p) for p in ps]
        except IndexError:
            print('In pron/word: {0}/{1}'.format(ps, w))
            raise
        # os = [['_W_'] + strip_consonants(strip_stress(p)) + ['_/W_'] for p in ps]
        out.append(os)

    return out

##############################################################################

def pron_bleu(prons_1, prons_2):
    """Returns the maximum BLEU score of two pronunciations."""
    logging.debug('Prons 1: {0}'.format(prons_1))
    logging.debug('Prons 2: {0}'.format(prons_2))

    best_score = -1.0
    best_p1 = None
    best_p2 = None
    for p1 in prons_1:
        for p2 in prons_2:
            if p1 == p2:
                bleu = 1.0
            elif len(p1) == 0:
                logging.warn('Zero-length pron!')
                bleu = 0.0
            elif len(p2) == 0:
                logging.warn('Zero-length pron!')
                bleu = 0.0
            else:
                print(p1)
                print(p2)

                bleu = nltk.translate.bleu_score.sentence_bleu([p2], p1)
                print('WORKS')
            if bleu > best_score:
                best_p1 = p1
                best_p2 = p2
                best_score = bleu

    return best_p1, best_p2, best_score

def pron_suffix_overlap(prons_1, prons_2):
    best_score, best_p1, best_p2 = -1.0, None, None
    for p1 in prons_1:
        for p2 in prons_2:
            if p1 == p2:
                score = 1.0
            elif (len(p1) == 0) or (len(p2) == 0):
                score = 0.0
            else:
                hits = 0.0
                for a, b in zip(reversed(p1), reversed(p2)):
                    if a == b:
                        hits += 1.0
                score = (2 * hits) / (len(p1) + len(p2))
            if score > best_score:
                best_score, best_p1, best_p2 = score, p1, p2
    return best_p1, best_p2, best_score

##############################################################################

def word_rhyming_table(words, prondict=None, pair_score_fn=pron_bleu, term_weights=None,
                       do_strip_consonants=True, MAX_DISTANCE=16):

    if prondict is None:
        prondict = collections.defaultdict(list)
        for word, syl in nltk.corpus.cmudict.entries():
            prondict[word].append(syl)

    prons = get_prons(words, prondict=prondict, unknown_pron_client=Pronounce())
    prons = preprocess_prons(prons, words=words, do_strip_consonants=do_strip_consonants)

    logging.debug('Words: {0}\nProns: {1}'.format(words, prons))
    # print('Words: {0}\nProns: {1}'.format(words, prons))

    # Your code goes here
    pair_scores = {}
    pair_prons = {}
    for i1, (ps1, w1) in enumerate(itertools.izip(prons, words)):

        # Relevant window: either the first MAX_DISTANCE words, or the MAX_DISTANCE / 2 words either way,
        # or the last MAX_DISTANCE words.
        if i1 < (MAX_DISTANCE / 2):
            lower_lim, upper_lim = 0, min(MAX_DISTANCE, len(words))
        elif i1 > (len(words) - MAX_DISTANCE / 2):
            lower_lim, upper_lim = max(0, len(words) - MAX_DISTANCE), len(words)
        else:
            lower_lim, upper_lim = max(0, (i1 - MAX_DISTANCE / 2)), min((i1 + MAX_DISTANCE / 2), len(words))

        # print('Window for i1={0}: {1} -- {2}'.format(i1, lower_lim, upper_lim))

        for i2, (ps2, w2) in enumerate(itertools.izip(prons[lower_lim:upper_lim],
                                                      words[lower_lim:upper_lim])):
            if w1 == w2:
                continue
            if ((w1, w2) in pair_scores) or ((w2, w1) in pair_scores):
                continue
            if is_stopword(w1) and is_stopword(w2):
                continue
            if is_strong_stopword(w1) or is_strong_stopword(w2):
                continue

            p1, p2, score = pair_score_fn(ps1, ps2)

            # Weights of scores
            if term_weights is not None:
                weight1 = 1.0
                if w1 in term_weights:
                    weight1 = term_weights[w1]
                weight2 = 1.0
                if w2 in term_weights:
                    weight2 = term_weights[w2]
                mean_weight = numpy.sqrt(weight1 * weight2)
                score *= mean_weight

            pair_scores[(w1, w2)] = score
            pair_prons[(w1, w2)] = (p1, p2)
            logging.debug('Scoring pair: {0} with {1:.3f}'.format((w1, w2), score))

    return pair_scores, pair_prons

##############################################################################

# Postprocessing - after the word scores have been computed

def rhyme_score_grid(words, prondict, pair_score_fn=pron_bleu, nonnegative=False,
                     do_strip_consonants=True):
    """Visualizes the pairwise rhyming scores."""
    pair_scores, pair_prons = word_rhyming_table(words, prondict=prondict, pair_score_fn=pair_score_fn,
                                                 do_strip_consonants=do_strip_consonants)

    score_grid = numpy.zeros((len(words), len(words))) - 1.0

    for i1, w1 in enumerate(words):
        for i2, w2 in enumerate(words):
            # check if it's there
            key = w1, w2
            if (w1, w2) not in pair_scores:
                if (w2, w1) not in pair_scores:
                    continue
                else:
                    key = w2, w1

            s = pair_scores[key]
            score_grid[i1, i2] = s
            score_grid[i2, i1] = s

    if nonnegative is True:
        score_grid[score_grid < 0.0] = 0.0

    return score_grid

def aggregate_score(grid):
    """Experiments with how to aggregate the score to get sensible results."""
    # Each word only needs to rhyme once.
    maxima = grid.max(axis=0)
    return numpy.average(maxima)

def binarize_grid(grid, thr=0.8):

    binary_grid = grid * 1
    binary_grid[grid <= thr] = 0
    binary_grid[grid > thr] = 1
    for i in xrange(binary_grid.shape[0]):
        binary_grid[i, i] = 1
    return binary_grid

def get_rhyme_graph(grid, words, thr=0.8):
    """Identify groups of words that look like they're on a grid.
    These are cliques in an adjacency graph.
    """
    word_keys = [w + '_{0}'.format(i) for i, w in enumerate(words)]
    binary_grid = binarize_grid(grid, thr=thr)
    G = networkx.Graph()
    G.add_nodes_from(word_keys)
    for i, x in enumerate(word_keys):
        for j, y in enumerate(word_keys):
            if binary_grid[i, j] > 0:
                G.add_edge(x, y)
    return G

def get_rhyme_groups(grid, words, thr=0.8):
    G = get_rhyme_graph(grid, words, thr=thr)
    cliques = list(networkx.enumerate_all_cliques(G))
    return cliques

def multi_clique_ratio(cliques, min_level=3):
    multi_cliques = [c for c in cliques if len(c) >= min_level]
    return multi_cliques

##############################################################################

def get_score(args_dict):
    """Call this from the judge. The args_dict should have the 'text' key:

    >>> args = {'text': ['This', 'is', 'some', 'rap', 'it', 'is', 'pretty', 'crap']}
    >>> score = get_score(args)

    """
    logging.info('Starting main...')
    _start_time = time.clock()

    cmudict = collections.defaultdict(list)
    for word, syl in nltk.corpus.cmudict.entries():
        cmudict[word].append(syl)

    words = tokenize(' '.join(args_dict['text']))

    # Your code goes here
    # pair_scores, pair_prons = word_rhyming_table(words, prondict=cmudict, pair_score_fn=pron_bleu)
    score_grid = rhyme_score_grid(words, prondict=cmudict, pair_score_fn=DEFAULT_PAIR_SCORE_FN)
    final_score = aggregate_score(score_grid)

    _end_time = time.clock()
    logging.info('[XXXX] done in {0:.3f} s'.format(_end_time - _start_time))

    return final_score

##############################################################################
DEFAULT_PAIR_SCORE_FN = pron_suffix_overlap

def build_argument_parser():
    parser = argparse.ArgumentParser(description=__doc__, add_help=True,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument('-t', '--text', action='store', nargs='+',
                        help='space-separated list of words.')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Turn on INFO messages.')
    parser.add_argument('--debug', action='store_true',
                        help='Turn on DEBUG messages.')

    return parser

def main(args):
    logging.info('Starting main...')
    _start_time = time.clock()

    cmudict = collections.defaultdict(list)
    for word, syl in nltk.corpus.cmudict.entries():
        cmudict[word].append(syl)

    words = tokenize(' '.join(args.text))

    # Your code goes here
    # pair_scores, pair_prons = word_rhyming_table(words, prondict=cmudict, pair_score_fn=pron_bleu)
    score_grid = rhyme_score_grid(words, prondict=cmudict, pair_score_fn=DEFAULT_PAIR_SCORE_FN)

    final_score = aggregate_score(score_grid)
    if __name__ == '__main__':
        print(final_score)

    _end_time = time.clock()
    logging.info('[XXXX] done in {0:.3f} s'.format(_end_time - _start_time))

    return final_score

##############################################################################

if __name__ == '__main__':
    parser = build_argument_parser()
    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.INFO)
    if args.debug:
        logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.DEBUG)

    main(args)

#back on my grid and all more potent than tylenol more focus on rhyming all the lyrics my mind can store

reason: import numpy,import networkx


text:
import numpy as np

yeo7_colors = np.array([[0, 0, 0, 255], [0, 118, 14, 255], [230, 148, 34, 255],
                        [205, 62, 78, 255], [120, 18, 134, 255],
                        [220, 248, 164, 255], [70, 130, 180, 255],
                        [196, 58, 250, 255]], dtype=np.uint8)

colormaps = {'yeo7': yeo7_colors}

max_stars_repo_name: josemariamoreira/BrainSpace

max_issues_repo_name: josemariamoreira/BrainSpace

max_forks_repo_name: josemariamoreira/BrainSpace


text:
"""Riemannian Line Search.

This module implements line search on Riemannian manifolds using geoopt.
This module uses the same syntax as a Torch optimizer
"""

from scipy.optimize.linesearch import scalar_search_wolfe2, scalar_search_armijo
import warnings
import torch
from .mixin import OptimMixin
from ..tensor import ManifoldParameter, ManifoldTensor
from ..manifolds import Euclidean

__all__ = ["RiemannianLineSearch"]

class LineSearchWarning(RuntimeWarning):
    pass

class RiemannianLineSearch(OptimMixin, torch.optim.Optimizer):
    r"""Riemannian line search optimizer.

    We try to minimize objective :math:`f\colon M\to \mathbb{R}`, in a search
    direction :math:`\eta`.
    This is done by minimizing the line search objective

    .. math::

        \phi(\alpha) = f(R_x(\alpha\eta)),

    where :math:`R_x` is the retraction at :math:`x`.
    Its derivative is given by

    .. math::

        \phi'(\alpha) = \langle\mathrm{grad} f(R_x(\alpha\eta)),\,
        \mathcal T_{\alpha\eta}(\eta) \rangle_{R_x(\alpha\eta)},

    where :math:`\mathcal T_\xi(\eta)` denotes the vector transport of :math:`\eta`
    to the point :math:`R_x(\xi)`.

    The search direction :math:`\eta` is defined recursively by

    .. math::

        \eta_{k+1} = -\mathrm{grad} f(R_{x_k}(\alpha_k\eta_k))
        + \beta \mathcal T_{\alpha_k\eta_k}(\eta_k)

    Here :math:`\beta` is the scale parameter. If :math:`\beta=0` this is steepest
    descent, other choices are Riemannian version of Fletcher-Reeves and
    Polak-Ribière scale parameters.

    Common conditions to accept the new point are the Armijo /
    sufficient decrease condition:

    .. math::

        \phi(\alpha)\leq \phi(0)+c_1\alpha\phi'(0)

    And additionally the curvature / (strong) Wolfe condition

    .. math::

        \phi'(\alpha)\geq c_2\phi'(0)

    The Wolfe conditions are more restrictive, but guarantee that search direction
    :math:`\eta` is a descent direction.

    The constants :math:`c_1` and :math:`c_2` satisfy :math:`c_1\in (0,1)`
    and :math:`c_2\in (c_1,1)`.

    Parameters
    ----------
    params : iterable
        iterable of parameters to optimize or dicts defining
        parameter groups
    line_search_method : ('wolfe', 'armijo', or callable)
        Which line_search_method to use. If callable it should be any method
        of signature `(phi, derphi, **kwargs) -> step_size`,
        where phi is scalar line search objective, and derphi is its derivative.
        If no suitable step size can be found, the method should return `None`.
        The following arguments are always passed in `**kwargs`:
        * **phi0:** float, Value of phi at 0
        * **old_phi0:** float, Value of phi at previous point
        * **derphi0:** float, Value derphi at 0
        * **old_derphi0:** float, Value of derphi at previous point
        * **old_step_size:** float, Stepsize at previous point
        If any of these arguments are undefined, they default to `None`.
        Additional arguments can be supplied through the `line_search_params` parameter
    line_search_params : dict
        Extra parameters to pass to `line_search_method`, for
        the parameters available to strong Wolfe see :meth:`strong_wolfe_line_search`.
        For Armijo backtracking parameters see :meth:`armijo_backtracking`.
    cg_method : ('steepest', 'fr', 'pr', or callable)
        Method used to compute the conjugate gradient scale parameter beta.
        If 'steepest', set the scale parameter to zero, which is equivalent
        to doing steepest descent. Use 'fr' for Fletcher-Reeves, or 'pr' for
        Polak-Ribière (NB: this setting requires an additional vector transport).
        If callable, it should be a function of signature
        `(params, states, **kwargs) -> beta`,
        where params are the parameters of this optimizer,
        states are the states associated to the parameters (self._states),
        and beta is a float giving the scale parameter. The keyword
        arguments are specified in optional parameter `cg_kwargs`.

    Other Paremeters
    ----------------
    compute_derphi : bool, optional
        If True, compute the derivative of the line search objective phi
        for every trial step_size alpha. If alpha is not zero, this requires
        a vector transport and an extra gradient computation. This is always set
        True if `line_search_method='wolfe'` and False if `'armijo'`, but needs
        to be manually set for a user implemented line search method.
    transport_grad : bool, optional
        If True, the transport of the gradient to the new point is computed
        at the end of every step. Set to `True` if Polak-Ribière is used, otherwise
        defaults to `False`.
    transport_search_direction: bool, optional
        If True, transport the search direction to new point at end of every step.
        Set to False if steepest descent is used, True Otherwise.
    fallback_stepsize : float
        fallback_stepsize to take if no point can be found satisfying
        line search conditions. See also :meth:`step` (default: 1)
    stabilize : int
        Stabilize parameters if they are off-manifold due to numerical
        reasons every `stabilize` steps (default: `None` -- no stabilize)
    cg_kwargs : dict
        Additional parameters to pass to the method used to compute the
        conjugate gradient scale parameter.

    Attributes
    ----------
    last_step_size : int or `None`
        Last step size taken. If `None` no suitable step size was
        found, and consequently no step was taken.
    step_size_history : List[int or `None`]
        List of all step sizes taken so far.
    line_search_method : callable
    line_search_params : dict
    cg_method : callable
    cg_kwargs : dict
    fallback_stepsize : float
    """

    def __init__(
        self,
        params,
        line_search_method="armijo",
        line_search_params=None,
        cg_method="steepest",
        cg_kwargs=None,
        compute_derphi=True,
        transport_grad=False,
        transport_search_direction=True,
        fallback_stepsize=1,
        stabilize=None,
    ):
        defaults = dict(
            line_search_method=line_search_method,
            line_search_params=line_search_params,
            cg_method=cg_method,
            cg_kwargs=cg_kwargs,
            compute_derphi=compute_derphi,
            transport_grad=transport_grad,
            transport_search_direction=transport_search_direction,
            fallback_stepsize=fallback_stepsize,
            stabilize=stabilize,
        )
        super(RiemannianLineSearch, self).__init__(
            params, defaults, stabilize=stabilize
        )
        self._params = []
        for group in self.param_groups:
            group.setdefault("step", 0)
            self._params.extend(group["params"])
        if len(self.param_groups) > 1:
            warning_string = """Multiple parameter groups detected.
            Line search parameters will be taken from first group.
            """
            warnings.warn(warning_string, UserWarning)

        self.compute_derphi = self.param_groups[0]["compute_derphi"]
        ls_method = self.param_groups[0]["line_search_method"]
        if ls_method == "wolfe":
            self.line_search_method = strong_wolfe_line_search
            self.compute_derphi = True
        elif ls_method == "armijo":
            self.line_search_method = armijo_backtracking
            self.compute_derphi = False
        elif callable(ls_method):
            self.line_search_method = ls_method
        else:
            raise ValueError(f"Unrecognized line search method '{ls_method}'")

        self.cg_kwargs = self.param_groups[0]["cg_kwargs"]
        if self.cg_kwargs is None:
            self.cg_kwargs = dict()
        self.transport_grad = self.param_groups[0]["transport_grad"]
        self.transport_search_direction = self.param_groups[0][
            "transport_search_direction"
        ]
        cg_method = self.param_groups[0]["cg_method"]
        if cg_method in ("steepest", "constant"):
            self.cg_method = cg_constant
            self.transport_search_direction = False
        elif cg_method in ("fr", "fletcher-reeves"):
            self.cg_method = cg_fletcher_reeves
        elif cg_method in ("pr", "polak-ribiere"):
            self.cg_method = cg_polak_ribiere
            self.transport_grad = True
        elif callable(cg_method):
            self.cg_method = cg_method
        else:
            raise ValueError(f"Unrecognized scale parameter method '{cg_method}'")

        self.line_search_params = self.param_groups[0]["line_search_params"]
        if self.line_search_params is None:
            self.line_search_params = dict()
        self.fallback_stepsize = self.param_groups[0]["fallback_stepsize"]

        self.old_phi0 = None
        self.old_derphi0 = None
        self.step_size_history = []
        self.last_step_size = None
        self._last_step = None
        self._grads_computed = False
        self.prev_loss = None
        self.closure = None
        self._step_size_dic = dict()

    def _phi(self, step_size):
        """Compute the line search objective, and store its derivatives in the state."""

        if step_size in self._step_size_dic:
            return self._step_size_dic[step_size]

        param_copy = [param.clone() for param in self._params]

        for point in self._params:
            state = self.state[point]
            if "search_direction" not in state:  # this shouldn't be possible actually
                raise ValueError("Search direction for parameter not computed.")
            if isinstance(point, (ManifoldParameter, ManifoldTensor)):
                manifold = point.manifold
            else:  # Use euclidean manifold
                manifold = Euclidean()

            search_direction = state["search_direction"]

            # compute retract and transport in search direction
            if self.compute_derphi:
                new_point, search_transported = manifold.retr_transp(
                    point, step_size * search_direction, search_direction
                )
                # This should not have any effect, but it does
                new_point = manifold.projx(new_point)
                state["search_transported"] = manifold.proju(
                    new_point, search_transported
                )
            else:
                new_point = manifold.retr(point, step_size * search_direction)

            with torch.no_grad():
                point.copy_(new_point)

        # recompute loss at new point
        phi = self.closure()

        if self.compute_derphi:
            # Store new gradients in state
            for point in self._params:
                grad = point.grad
                if grad is None:
                    continue
                if isinstance(point, (ManifoldParameter, ManifoldTensor)):
                    manifold = point.manifold
                else:  # Use euclidean manifold
                    manifold = Euclidean()

                state = self.state[point]

                # project gradient onto tangent space
                grad = manifold.egrad2rgrad(point, grad)

                state["new_grad"] = grad
                state["der_phi"] = torch.sum(
                    manifold.inner(point, grad, state["search_transported"])
                ).item()

            self._grads_computed = True

        # roll back parameters to before step, save new point is state
        with torch.no_grad():
            for point, old_point in zip(self._params, param_copy):
                state = self.state[point]
                state["new_point"] = point.clone()

                point.copy_(old_point)

        self._step_size_dic[step_size] = phi
        self._last_step = step_size
        return phi

    def _derphi(self, step_size):
        """Compute derivative of phi.

        The derivative of phi is given by computing inner
        product between all tensor gradients at target point and those at source point.
        The source gradients are transported to the target point, and both gradients are
        projected.
        """

        if not self.compute_derphi:
            raise ValueError("Cannot call _derphi if self.compute_derphi=False!")

        # Call _phi to compute gradients; Does nothing if _phi was
        # already called with this stepsize during this step
        self._phi(step_size)

        derphi = 0
        for point in self._params:
            state = self.state[point]
            if "der_phi" not in state:
                continue

            derphi += state["der_phi"]

        return derphi

    def _init_loss(self, recompute_gradients=False):
        """Compute loss, gradients and search direction at start of line search.

        Use `recompute_gradients=True` if gradients have changed between
        consecutive calls of `step`.
        """

        if recompute_gradients or (not self._grads_computed):
            loss = self.closure()
            reuse_grads = False
        else:
            loss = self.prev_loss
            reuse_grads = True

        self._step_size_dic = dict()

        for point in self._params:
            state = self.state[point]
            if isinstance(point, (ManifoldParameter, ManifoldTensor)):
                manifold = point.manifold
            else:  # Use euclidean manifold
                manifold = Euclidean()
            if reuse_grads:
                grad = state["new_grad"]
            else:
                grad = point.grad
                grad = manifold.egrad2rgrad(point, grad)

            if "grad" in state:
                state["prev_grad"] = state["grad"]
                state["prev_grad_norm_squared"] = torch.sum(
                    manifold.inner(point, state["grad"])
                ).item()
            state["grad"] = grad

        derphi0 = self._compute_search_direction()

        self._grads_computed = True

        return loss, derphi0

    def _compute_search_direction(self):
        """Compute the search direction.

        If the direction is not a descent direction, revert to steepest descent.
        """

        first_time = False
        for point in self._params:
            state = self.state[point]
            if "search_direction" not in state:
                state["search_direction"] = -state["grad"]
                first_time = True

        if not first_time:
            beta = self.cg_method(self._params, self.state, **self.cg_kwargs)

            for point in self._params:
                state = self.state[point]

                if beta != 0:
                    state["search_direction"] = (
                        -state["grad"] + beta * state["search_transported"]
                    )
                else:
                    state["search_direction"] = -state["grad"]

        # Deriphative of phi at zero is inner product grad and search direction
        derphi0 = 0
        for point in self._params:
            state = self.state[point]
            if isinstance(point, (ManifoldParameter, ManifoldTensor)):
                manifold = point.manifold
            else:  # Use euclidean manifold
                manifold = Euclidean()
            derphi0 += torch.sum(
                manifold.inner(point, state["grad"], state["search_direction"])
            ).item()

        # If search direction is not a descent direction, revert to gradient
        if derphi0 >= 0:
            derphi0 = 0
            for point in self._params:
                state = self.state[point]
                if isinstance(point, (ManifoldParameter, ManifoldTensor)):
                    manifold = point.manifold
                else:  # Use euclidean manifold
                    manifold = Euclidean()
                derphi0 -= torch.sum(manifold.inner(point, state["grad"])).item()
                state["search_direction"] = -state["grad"]

        return derphi0

    def step(self, closure, force_step=False, recompute_gradients=False, no_step=False):
        """Do a linesearch step.

        Parameters
        ----------
        closure : callable
            A closure that reevaluates the model and returns the loss.
        force_step : bool (optional)
            If `True`, take a unit step of size `self.fallback_stepsize`
            if no suitable step size can be found.
            If `False`, no step is taken in this situation. (default: `False`)
        recompute_gradients : bool (optional)
            If True, recompute the gradients. Use this if the parameters
            have changed in between consecutive steps. (default: False)
        no_step : bool (optional)
            If True, just compute step size and do not perform the step.
            (default: False)
        """

        self.closure = closure
        phi0, derphi0 = self._init_loss(recompute_gradients=recompute_gradients)
        self._step_size_dic = dict()

        phi_information = {
            "phi0": phi0,
            "derphi0": derphi0,
            "old_phi0": self.old_phi0,
            "old_derphi0": self.old_derphi0,
            "old_step_size": self.last_step_size,
        }
        step_size = self.line_search_method(
            self._phi, self._derphi, **phi_information, **self.line_search_params
        )

        self.step_size_history.append(step_size)
        self.last_step_size = step_size

        self.old_phi0 = phi0
        self.old_derphi0 = derphi0

        # Ensure that the last step for which we computed the closure coincides with
        # proposed step size, so that we can reuse the gradients and retract.
        # This is very rare, and should only happen if force_step=True and no stepsize
        # was found, or for user-defined linesearch methods.
        if self._last_step != step_size or not self.compute_derphi:
            self._grads_computed = False
        redo_retract = self._last_step != step_size

        # If it fails to find a good step, probably we have convergence
        if step_size is None:
            if force_step:
                step_size = self.fallback_stepsize
                self._grads_computed = False
            elif (
                self.last_step_size is None
            ):  # Warn if step_size is None twice in a row
                warning_string = """No suitable step size could be found, and no step
                was taken. Call `step` with `force_step=True` to take a step anyway.
                """
                warnings.warn(warning_string, LineSearchWarning)

        for point in self._params:
            if step_size is None or no_step:
                continue
            state = self.state[point]
            if "search_direction" not in state:
                continue
            if isinstance(point, (ManifoldParameter, ManifoldTensor)):
                manifold = point.manifold
            else:  # Use euclidean manifold
                manifold = Euclidean()

            search_direction = state["search_direction"]

            # Compute retract if suggested step size is not the last one we tried (rare)
            if redo_retract:
                new_point = manifold.retr(point, step_size * search_direction)
            else:
                new_point = state["new_point"]

            # Use retract to perform the step, and transport the search direction
            if self.transport_search_direction:
                search_transported = manifold.transp_follow_retr(
                    point, step_size * search_direction, search_direction
                )
                state["search_transported"] = search_transported

            if self.transport_grad:
                grad_transport = manifold.transp_follow_retr(
                    point, step_size * search_direction, state["grad"]
                )
                state["grad_transported"] = grad_transport

            with torch.no_grad():  # Take suggested step
                point.copy_(new_point)

        for group in self.param_groups:
            group["step"] += 1
            if (
                group["stabilize"] is not None
                and group["step"] % group["stabilize"] == 0
            ):
                self.stabilize_group(group)

        # Update loss value
        if step_size is not None:
            new_loss = self._phi(step_size)
            self.prev_loss = new_loss
        else:
            new_loss = self.prev_loss
        return new_loss

    def stabilize_group(self, group):
        for p in group["params"]:
            if not isinstance(p, (ManifoldParameter, ManifoldTensor)):
                continue
            state = self.state[p]
            if not state:  # due to None grads
                continue
            manifold = p.manifold
            p.copy_(manifold.projx(p))

#################################################################################
# Conjugate gradient scale factor
#################################################################################

def cg_constant(params, states, alpha=0, **kwargs):
    """Constant scale parameter. If alpha=0, then this is steepest descent."""

    return alpha

def cg_fletcher_reeves(params, states, **kwargs):
    r"""Fletcher-Reeves scale parameter.

    This is given by

    .. math::
        \beta_{k+1}^{FR} = \frac{\langle\nabla f(x_{k+1},\,
        \nabla f(x_{k+1}\rangle_{x_{k+1}}
        {\langle\nabla f(x_k),\nabla f(x_k)\rangle_{x_k}}
    """
    numerator = 0
    denominator = 0
    for point in params:
        state = states[point]

        # Can't compute beta, probably first step hasn't been taken yet
        if "prev_grad_norm_squared" not in state:
            return 0

        if isinstance(point, (ManifoldParameter, ManifoldTensor)):
            manifold = point.manifold
        else:  # Use euclidean manifold
            manifold = Euclidean()
        numerator += torch.sum(manifold.inner(point, state["grad"])).item()
        denominator += state["prev_grad_norm_squared"]

    if denominator == 0:
        return 0
    else:
        return numerator / denominator

def cg_polak_ribiere(params, states, **kwargs):
    r"""Polak-Ribière scale parameter.

    This is given by

    .. math::
        \beta_{k+1}^{PR} = \frac{\langle\nabla f(x_{k+1}
        ,\,\nabla f(x_{k+1})-\mathcal T_{\alpha_k\eta_k}\nabla f(x_k)\rangle_{x_{k+1}}}
        {\langle\nabla f(x_k),\,\nabla f(x_k)\rangle_{x_k}}
    """

    numerator = 0
    denominator = 0
    for point in params:
        state = states[point]

        # Can't compute beta, probably first step hasn't been taken yet.
        if "grad_transported" not in state:
            return 0

        if isinstance(point, (ManifoldParameter, ManifoldTensor)):
            manifold = point.manifold
        else:  # Use euclidean manifold
            manifold = Euclidean()
        numerator += torch.sum(
            manifold.inner(
                point, state["grad"], state["grad"] - state["grad_transported"]
            )
        ).item()
        denominator += state["prev_grad_norm_squared"]

    if denominator == 0:
        return 0
    else:
        return numerator / denominator

#################################################################################
# Line search methods
#################################################################################

def strong_wolfe_line_search(
    phi,
    derphi,
    phi0=None,
    old_phi0=None,
    derphi0=None,
    c1=1e-4,
    c2=0.9,
    amax=None,
    **kwargs,
):
    """
    Scalar line search method to find step size satisfying strong Wolfe conditions.

    Parameters
    ----------
    c1 : float, optional
        Parameter for Armijo condition rule.
    c2 : float, optional
        Parameter for curvature condition rule.
    amax : float, optional
        Maximum step size

    Returns
    -------
    step_size : float
        The next step size
    """

    step_size, _, _, _ = scalar_search_wolfe2(
        phi,
        derphi,
        phi0=phi0,
        old_phi0=old_phi0,
        c1=c1,
        c2=c2,
        amax=amax,
    )

    return step_size

def armijo_backtracking(
    phi,
    derphi,
    phi0=None,
    derphi0=None,
    old_phi0=None,
    c1=1e-4,
    amin=0,
    amax=None,
    **kwargs,
):
    """Scalar line search method to find step size satisfying Armijo conditions.

    Parameters
    ----------
    c1 : float, optional
        Parameter for Armijo condition rule.
    amax, amin : float, optional
        Maxmimum and minimum step size
    """

    # TODO: Allow different schemes to choose initial step size

    if old_phi0 is not None and derphi0 != 0:
        alpha0 = 1.01 * 2 * (phi0 - old_phi0) / derphi0
    else:
        alpha0 = 1.0
    if alpha0 <= 0:
        alpha0 = 1.0
    if amax is not None:
        alpha0 = min(alpha0, amax)

    step_size, _ = scalar_search_armijo(
        phi, phi0, derphi0, c1=c1, alpha0=alpha0, amin=amin
    )

    return step_size

max_stars_repo_name: leonMatzner/geoopt

max_issues_repo_name: leonMatzner/geoopt

max_forks_repo_name: leonMatzner/geoopt

reason: from scipy