96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
from django.core.management.base import BaseCommand
|
|
from django.contrib.auth.models import User
|
|
|
|
from student.models import UserTestGroup
|
|
|
|
import random
|
|
import sys
|
|
import datetime
|
|
|
|
import json
|
|
from pytz import UTC
|
|
|
|
|
|
def group_from_value(groups, v):
|
|
''' Given group: (('a',0.3),('b',0.4),('c',0.3)) And random value
|
|
in [0,1], return the associated group (in the above case, return
|
|
'a' if v<0.3, 'b' if 0.3<=v<0.7, and 'c' if v>0.7
|
|
'''
|
|
sum = 0
|
|
for (g, p) in groups:
|
|
sum = sum + p
|
|
if sum > v:
|
|
return g
|
|
return g # For round-off errors
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = \
|
|
''' Assign users to test groups. Takes a list
|
|
of groups:
|
|
a:0.3,b:0.4,c:0.3 file.txt "Testing something"
|
|
Will assign each user to group a, b, or c with
|
|
probability 0.3, 0.4, 0.3. Probabilities must
|
|
add up to 1.
|
|
|
|
Will log what happened to file.txt.
|
|
'''
|
|
|
|
def handle(self, *args, **options):
|
|
if len(args) != 3:
|
|
print "Invalid number of options"
|
|
sys.exit(-1)
|
|
|
|
# Extract groups from string
|
|
group_strs = [x.split(':') for x in args[0].split(',')]
|
|
groups = [(group, float(value)) for group, value in group_strs]
|
|
print "Groups", groups
|
|
|
|
## Confirm group probabilities add up to 1
|
|
total = sum(zip(*groups)[1])
|
|
print "Total:", total
|
|
if abs(total - 1) > 0.01:
|
|
print "Total not 1"
|
|
sys.exit(-1)
|
|
|
|
## Confirm groups don't already exist
|
|
for group in dict(groups):
|
|
if UserTestGroup.objects.filter(name=group).count() != 0:
|
|
print group, "already exists!"
|
|
sys.exit(-1)
|
|
|
|
group_objects = {}
|
|
|
|
f = open(args[1], "a+")
|
|
|
|
## Create groups
|
|
for group in dict(groups):
|
|
utg = UserTestGroup()
|
|
utg.name = group
|
|
utg.description = json.dumps({"description": args[2]},
|
|
{"time": datetime.datetime.now(UTC).isoformat()})
|
|
group_objects[group] = utg
|
|
group_objects[group].save()
|
|
|
|
## Assign groups
|
|
users = list(User.objects.all())
|
|
count = 0
|
|
for user in users:
|
|
if count % 1000 == 0:
|
|
print count
|
|
count = count + 1
|
|
v = random.uniform(0, 1)
|
|
group = group_from_value(groups, v)
|
|
group_objects[group].users.add(user)
|
|
f.write("Assigned user {name} ({id}) to {group}\n".format(name=user.username,
|
|
id=user.id,
|
|
group=group))
|
|
|
|
## Save groups
|
|
for group in group_objects:
|
|
group_objects[group].save()
|
|
f.close()
|
|
|
|
# python manage.py assigngroups summary_test:0.3,skip_summary_test:0.7 log.txt "Do previews of future materials help?"
|
|
# python manage.py assigngroups skip_capacitor:0.3,capacitor:0.7 log.txt "Do we show capacitor in linearity tutorial?"
|