<commit_before>from __future__ import division
import numpy as np
from pyomo.environ import *
from pyomo.opt import SolverFactory

from math import sqrt
from hopper import Hopper
from hopperUtil import *

desiredPrecision = 2
N = 20
tf = 2*1.6
legLength = 0.18
r0 = [0, legLength/2]
rf = [6*legLength, legLength]
v0 = [0, 0]
w0 = 0
hipOffset = {'front': {'x': 0.5, 'z': -0.25}, 'hind': {'x': -0.5, 'z': -0.25}}

matlab_hopper = eng.Hopper(legLength, hipOffset)
hop = Hopper(N, eng, matlab_hopper)
# hop.mdt_precision = int(ceil(-np.log2(desiredPrecision)))
hop.dtBounds = tuple((1/sqrt(legLength/9.81))*np.array([0.01, 0.2]))
hop.dtNom = 0.1*(1/sqrt(legLength/9.81))
hop.rotationMax = np.pi/8
hop.nOrientationSectors = 1 #int(floor(np.pi/8/desiredPrecision))
print 'hop.nOrientationSectors = %d' % hop.nOrientationSectors
hop.velocityMax = 3.
hop.positionMax = 1.5*rf[0]/legLength
hop.forceMax = 3.
#addThreePlatfomWorld(hop, legLength)
addFlatWorld(hop, legLength)
hop.constructVisualizer()
m_nlp = hop.constructPyomoModel()

m_nlp.forceSlacks = Var(m_nlp.feet, m_nlp.R2_INDEX, m_nlp.t, bounds=(0.0, hop.forceMax))
#m_nlp.hipTorqueSlacks = Var(m_nlp.feet, m_nlp.t, bounds=(0.0, hop.forceMax))

def _forceSlackRuleUB(m, foot, xz, t):
    return m.f[foot, xz, t] <= m.forceSlacks[foot, xz, t]
m_nlp.forceSlackUBConstraint = Constraint(m_nlp.feet, m_nlp.R2_INDEX, m_nlp.t, rule=_forceSlackRuleUB)

def _forceSlackRuleLB(m, foot, xz, t):
    return m.f[foot, xz, t] >= -m.forceSlacks[foot, xz, t]
m_nlp.forceSlackLBConstraint = Constraint(m_nlp.feet, m_nlp.R2_INDEX, m_nlp.t, rule=_forceSlackRuleLB)

def objRule(m):
    #     return sum(m.beta[foot, bv, ti]**2 for foot in m.feet for bv in m.BV_INDEX for ti in m.t)
    #     + sum(m.pdd[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t)
    #return sum(m.f[foot, i, j]**2 + m.pd[foot, i, j]**2 + m.pdd[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t) + sum(m.T[ti]**2 for ti in m.t)
    #return sum(m.f[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t) + sum(m.hipTorque[foot, ti]**2 for foot in m.feet for ti in m.t)
    #return sum(m.f[foot, i, j]**2 + m.pd[foot, i, j]**2  + m.pdd[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t) + sum(m.hipTorque[foot, ti]**2 for foot in m.feet for ti in m.t) + summation(m.dt)

    return sum(m.f[foot, i, j]**2 + m.pdd[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t) + sum(m.hipTorque[foot, ti]**2 for foot in m.feet for ti in m.t) + summation(m.dt)

    #return summation(m.forceSlacks)

m_nlp.Obj = Objective(rule=objRule, sense=minimize)

m_nlp.rx0 = Constraint(expr=m_nlp.r['x',m_nlp.t[1]] == r0[0]/legLength)
#m_nlp.rz0 = Constraint(expr=m_nlp.r['z',m_nlp.t[1]] <= 1)

m_nlp.th0 = Constraint(expr=m_nlp.th[m_nlp.t[1]] == 0)
<commit_msg>Add helper functions for L1, L2, and Linf norms<commit_after>from __future__ import division
import numpy as np
from pyomo.environ import *
from pyomo.opt import SolverFactory

from math import sqrt
from hopper import Hopper
from hopperUtil import *

desiredPrecision = 2
N = 20
tf = 2*1.6
legLength = 0.18
r0 = [0, legLength/2]
rf = [6*legLength, legLength]
v0 = [0, 0]
w0 = 0
hipOffset = {'front': {'x': 0.5, 'z': -0.25}, 'hind': {'x': -0.5, 'z': -0.25}}

matlab_hopper = eng.Hopper(legLength, hipOffset)
hop = Hopper(N, eng, matlab_hopper)
# hop.mdt_precision = int(ceil(-np.log2(desiredPrecision)))
hop.dtBounds = tuple((1/sqrt(legLength/9.81))*np.array([0.01, 0.2]))
hop.dtNom = 0.1*(1/sqrt(legLength/9.81))
hop.rotationMax = np.pi/8
hop.nOrientationSectors = 1 #int(floor(np.pi/8/desiredPrecision))
print 'hop.nOrientationSectors = %d' % hop.nOrientationSectors
hop.velocityMax = 3.
hop.positionMax = 1.5*rf[0]/legLength
hop.forceMax = 3.
#addThreePlatfomWorld(hop, legLength)
addFlatWorld(hop, legLength)
hop.constructVisualizer()
m_nlp = hop.constructPyomoModel()
def normL2(m, var):
    index = var.index_set()
    return sum(var[i]**2 for i in index)

def normL1(m, var):
    index = var.index_set()
    slackName = '%sSlacks' % var.cname()
    lbName = '%sLB' % slackName
    ubName = '%sUB' % slackName
    slackMax = max([max(np.abs(var[i].bounds)) for i in index])
    setattr(m, slackName, Var(index, bounds=(0.0, slackMax)))

    def _lbRule(m, *args):
        return var[args] <= getattr(m, slackName)[args]
    setattr(m, lbName, Constraint(index, rule=_lbRule))

    def _ubRule(m, *args):
        return var[args] >= -getattr(m, slackName)[args]
    setattr(m, ubName, Constraint(index, rule=_ubRule))

    return summation(getattr(m, slackName))

def normLInfinity(m, var):
    index = var.index_set()
    slackName = '%sSlack' % var.cname()
    lbName = '%sLB' % slackName
    ubName = '%sUB' % slackName
    slackMax = max([max(np.abs(var[i].bounds)) for i in index])
    setattr(m, slackName, Var(bounds=(0.0, slackMax)))

    def _lbRule(m, *args):
        return var[args] <= getattr(m, slackName)
    setattr(m, lbName, Constraint(index, rule=_lbRule))

    def _ubRule(m, *args):
        return var[args] >= -getattr(m, slackName)
    setattr(m, ubName, Constraint(index, rule=_ubRule))

    return getattr(m, slackName)


norm = normLInfinity;

def objRule(m):
    #     return sum(m.beta[foot, bv, ti]**2 for foot in m.feet for bv in m.BV_INDEX for ti in m.t)
    #     + sum(m.pdd[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t)
    #return sum(m.f[foot, i, j]**2 + m.pd[foot, i, j]**2 + m.pdd[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t) + sum(m.T[ti]**2 for ti in m.t)
    #return sum(m.f[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t) + sum(m.hipTorque[foot, ti]**2 for foot in m.feet for ti in m.t)
    #return sum(m.f[foot, i, j]**2 + m.pd[foot, i, j]**2  + m.pdd[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t) + sum(m.hipTorque[foot, ti]**2 for foot in m.feet for ti in m.t) + summation(m.dt)

    #return sum(m.f[foot, i, j]**2 + m.pdd[foot, i, j]**2 for foot in m.feet for i in m.R2_INDEX for j in m.t) + sum(m.hipTorque[foot, ti]**2 for foot in m.feet for ti in m.t) + summation(m.dt)

    return norm(m, m.f) + norm(m, m.pdd) + norm(m, m.hipTorque) + summation(m.dt)

m_nlp.Obj = Objective(rule=objRule, sense=minimize)

m_nlp.rx0 = Constraint(expr=m_nlp.r['x',m_nlp.t[1]] == r0[0]/legLength)
#m_nlp.rz0 = Constraint(expr=m_nlp.r['z',m_nlp.t[1]] <= 1)

m_nlp.th0 = Constraint(expr=m_nlp.th[m_nlp.t[1]] == 0)

{'fraction_non_alphanumeric': 0.17157014555589967, 'fraction_numerical': 0.027872406317745433, 'mean_word_length': 4.022550544323484, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4519198', 'n_tokens_mistral': 2930, 'n_tokens_neox': 2796, 'n_words': 681}